diff --git a/agent.py b/agent.py index 7f98482..029650e 100644 --- a/agent.py +++ b/agent.py @@ -1,51 +1,296 @@ """ -agent.py — AI-агент с RAG-памятью на базе Ollama (llama3) и Qdrant. +rag_agent.py — AI-агент с RAG-памятью на базе Ollama + Qdrant. -Использует create_agent из langchain.agents с системным промптом, -инструктирующим агента работать с базой знаний. +Стек: 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 langchain_ollama import ChatOllama -from langchain_core.messages import SystemMessage +from qdrant_client import QdrantClient as _QdrantClient +from qdrant_client.http.models import Distance, VectorParams -from tools import search_knowledge_base, add_to_knowledge_base +# --------------------------------------------------------------------------- +# Конфигурация +# --------------------------------------------------------------------------- -# ── Модель ──────────────────────────────────────────────────────────────────── -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-эмбеддинги) +# --------------------------------------------------------------------------- + +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. Если в базе найдена релевантная информация — используй её как основу ответа - и явно ссылайся на источник. +1. ВСЕГДА начинай с поиска через search_knowledge_base перед ответом на вопрос. +2. Если найдена релевантная информация — используй её и ссылайся на источник. 3. Если пользователь просит сохранить информацию — используй add_to_knowledge_base. -4. Если база знаний не содержит нужной информации — честно сообщи об этом - и постарайся помочь на основе своих знаний. +4. Если база знаний пуста по теме — честно сообщи об этом и помоги по своим знаниям. 5. Отвечай на том же языке, на котором задан вопрос. """ -# ── Агент ───────────────────────────────────────────────────────────────────── -rag_tools = [search_knowledge_base, add_to_knowledge_base] - agent = create_agent( model=llm, - tools=rag_tools, + 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 "Агент не вернул ответ." \ No newline at end of file + 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]") \ No newline at end of file