Добавить rag_agent.py

rag_agent.py

@@ -0,0 +1,296 @@
+"""
+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 langgraph.prebuilt import create_react_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_react_agent из langgraph — актуальный API v0.3+)
+# ---------------------------------------------------------------------------
+
+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_react_agent(
+    model=llm,
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    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]")