Агент с RAG‑памятью: rag_agent.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py

@@ -0,0 +1,208 @@
+<|channel|>final code<|message|>import os
+from pathlib import Path
+from typing import List
+
+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, Tool
+from langchain.schema import Document
+
+# --------------------------------------------------------------------------- #
+# Конфигурация подключения к Qdrant и 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 = "nomic-embed-text"
+LLM_MODEL = "llama3"
+
+# --------------------------------------------------------------------------- #
+# Векторное хранилище
+# --------------------------------------------------------------------------- #
+
+class RagVectorStore:
+    """Обёртка над QdrantVectorStore с автоматическим эмбеддингом Ollama."""
+
+    def __init__(self, collection_name: str = COLLECTION_NAME):
+        self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+        self.store = QdrantVectorStore(
+            url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
+            collection_name=collection_name,
+            embeddings=self.embeddings,
+        )
+        # Создаём коллекцию, если её ещё нет
+        if not self.store.client.get_collection(collection_name).exists:
+            self.store.create_collection()
+
+    def add_documents(self, docs: List[Document]) -> None:
+        """Добавляет список документов в хранилище."""
+        self.store.add_documents(docs)
+
+    def search(
+        self,
+        query: str,
+        k: int = 5,
+    ) -> List[tuple[str, float]]:
+        """Семантический поиск по запросу. Возвращает кортежи (текст, релевантность)."""
+        results = self.store.similarity_search_with_score(query, k=k)
+        return [(doc.page_content, score) for doc, score in results]
+
+# --------------------------------------------------------------------------- #
+# Чанкинг документов
+# --------------------------------------------------------------------------- #
+
+def chunk_document(text: str, title: str) -> List[Document]:
+    """Разбивает текст на чанки и возвращает список Document."""
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+    )
+    chunks = splitter.split_text(text)
+    return [
+        Document(page_content=chunk, metadata={"title": title})
+        for chunk in chunks
+    ]
+
+# --------------------------------------------------------------------------- #
+# Инструменты для агента
+# --------------------------------------------------------------------------- #
+
+@tool("search_knowledge_base")
+def search_knowledge_base(
+    query: str,
+    max_results: int = 5,
+) -> List[str]:
+    """
+    Семантический поиск в базе знаний.
+
+    Parameters:
+        query (str): поисковый запрос.
+        max_results (int): максимальное количество результатов.
+
+    Returns:
+        list of strings: найденные тексты с оценкой релевантности.
+    """
+    results = rag_store.search(query, k=max_results)
+    return [f"{text} (score={score:.4f})" for text, score in results]
+
+@tool("add_to_knowledge_base")
+def add_to_knowledge_base(
+    content: str,
+    title: str,
+) -> str:
+    """
+    Добавляет новый документ в базу знаний.
+
+    Parameters:
+        content (str): полный текст документа.
+        title (str): заголовок/название документа.
+
+    Returns:
+        str: сообщение об успехе.
+    """
+    docs = chunk_document(content, title)
+    rag_store.add_documents(docs)
+    return f"Документ '{title}' успешно добавлен в базу знаний."
+
+# --------------------------------------------------------------------------- #
+# Инициализация хранилища
+# --------------------------------------------------------------------------- #
+
+rag_store = RagVectorStore()
+
+# --------------------------------------------------------------------------- #
+# Создание агента с RAG‑интеграцией
+# --------------------------------------------------------------------------- #
+
+llm = OllamaLLM(model=LLM_MODEL)
+
+tools: List[Tool] = [
+    Tool.from_function(
+        func=search_knowledge_base,
+        name="search_knowledge_base",
+        description="Используйте для поиска информации в базе знаний.",
+    ),
+    Tool.from_function(
+        func=add_to_knowledge_base,
+        name="add_to_knowledge_base",
+        description="Добавьте новый документ в базу знаний.",
+    ),
+]
+
+agent = create_agent(
+    llm=llm,
+    tools=tools,
+    system_message=(
+        "Вы — интеллектуальный агент, способный искать и добавлять информацию "
+        "в локальную базу знаний. При необходимости используйте инструменты "
+        "'search_knowledge_base' и 'add_to_knowledge_base'."
+    ),
+)
+
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+
+# --------------------------------------------------------------------------- #
+# Клиент для загрузки документов из директории
+# --------------------------------------------------------------------------- #
+
+def load_documents_from_directory(directory: str) -> None:
+    """Загружает все текстовые файлы из указанной папки в базу знаний."""
+    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")
+            title = file_path.stem
+            add_to_knowledge_base(content=text, title=title)
+
+# --------------------------------------------------------------------------- #
+# Интерактивный клиент
+# --------------------------------------------------------------------------- #
+
+def interactive_cli():
+    """Простой CLI для взаимодействия с агентом."""
+    print("=== RAG Агент ===")
+    print("Команды: /add <path>, /search <query>, /quit")
+    while True:
+        try:
+            user_input = input("> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            break
+
+        if not user_input:
+            continue
+
+        if user_input.lower() == "/quit":
+            print("Выход.")
+            break
+        elif user_input.startswith("/add"):
+            _, *path_parts = user_input.split(maxsplit=1)
+            if path_parts:
+                file_path = Path(path_parts[0])
+                if file_path.is_file():
+                    content = file_path.read_text(encoding="utf-8")
+                    title = file_path.stem
+                    print(add_to_knowledge_base(content, title))
+                else:
+                    print(f"Файл не найден: {file_path}")
+            else:
+                print("Укажите путь к файлу.")
+        elif user_input.startswith("/search"):
+            _, *query_parts = user_input.split(maxsplit=1)
+            if query_parts:
+                query = query_parts[0]
+                results = search_knowledge_base(query, max_results=3)
+                for r in results:
+                    print(r)
+            else:
+                print("Укажите поисковый запрос.")
+        else:
+            # Любой другой ввод передаём агенту
+            response = agent_executor.invoke({"input": user_input})
+            print(response["output"])
+
+if __name__ == "__main__":
+    interactive_cli()