# main.py # Полностью рабочий пример агента с RAG‑памятью на базе Qdrant и Ollama. # Весь код написан в одном файле для простоты демонстрации. # # Требования: # - Python 3.10+ # - Qdrant (работает по умолчанию на localhost:6333) # - Ollama (llama3 + nomic-embed-text) # - LangChain 1.2.10+ и связанные пакеты # # Запуск: # 1. Убедитесь, что Qdrant и Ollama запущены. # 2. pip install -r requirements.txt # 3. python main.py # # После запуска появится интерактивный клиент с командами: # /add – добавить документ в базу знаний. # /search – выполнить поиск. # /quit – выйти. # Любой другой ввод – агент обработает как обычный запрос. import os import sys import textwrap from pathlib import Path from typing import List, Dict, Any from langchain_ollama import ChatOllama, OllamaEmbeddings from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain.tools import tool from langchain.agents import create_agent, AgentExecutor, AgentType from langchain_core.messages import HumanMessage # --------------------------------------------------------------------------- # Конфигурация # --------------------------------------------------------------------------- QDRANT_URL = "http://localhost:6333" COLLECTION_NAME = "knowledge" EMBEDDING_MODEL = "nomic-embed-text" LLM_MODEL = "llama3" # --------------------------------------------------------------------------- # Векторное хранилище # --------------------------------------------------------------------------- # Инициализируем эмбеддер и клиент Qdrant embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) vector_store = QdrantVectorStore( url=QDRANT_URL, collection_name=COLLECTION_NAME, embedding=embeddings, ) # --------------------------------------------------------------------------- # Чанкинг # --------------------------------------------------------------------------- text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) # --------------------------------------------------------------------------- # Инструменты # --------------------------------------------------------------------------- @tool("search_knowledge_base", "Semantic search in the knowledge base") def search_knowledge_base(query: str, max_results: int = 5) -> str: """Возвращает топ‑N релевантных фрагментов из Qdrant.""" results = vector_store.similarity_search_with_score(query, k=max_results) if not results: return "No relevant documents found." # Формируем читаемый ответ formatted = [] for i, (doc, score) in enumerate(results, 1): formatted.append(f"{i}. (score={score:.3f})\n{doc.page_content[:500]}\n---") return "\n".join(formatted) @tool("add_to_knowledge_base", "Add a document to the knowledge base") def add_to_knowledge_base(content: str, title: str = "") -> str: """Разбивает документ на чанки, эмбеддит и сохраняет в Qdrant.""" # Разбиваем на чанки chunks = text_splitter.split_text(content) # Создаём Document объекты с метаданными from langchain.schema import Document docs = [Document(page_content=chunk, metadata={"title": title, "source": title}) for chunk in chunks] # Добавляем в хранилище vector_store.add_documents(docs) return f"Added {len(docs)} chunks from '{title}'." # --------------------------------------------------------------------------- # Агент # --------------------------------------------------------------------------- # Системный промпт, который подсказывает агенту использовать инструменты SYSTEM_PROMPT = textwrap.dedent(""" You are an intelligent assistant with access to a knowledge base. Use the following tools when you need to retrieve or store information: - search_knowledge_base: Perform a semantic search in the knowledge base. - add_to_knowledge_base: Add new content to the knowledge base. When you answer a user query, first decide if you need to search the base. If you do, call search_knowledge_base with a relevant query. If you need to add new information, call add_to_knowledge_base. Otherwise, answer directly using your internal knowledge. """) # Создаём LLM llm = ChatOllama(model=LLM_MODEL, temperature=0.2) # Создаём агент agent = create_agent( llm=llm, tools=[search_knowledge_base, add_to_knowledge_base], system_message=SYSTEM_PROMPT, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, ) # Обёртка для удобного вызова agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True) # --------------------------------------------------------------------------- # Загрузка документов из директории (инициализация) # --------------------------------------------------------------------------- def load_documents_from_dir(directory: str) -> None: """Загружает все .txt и .md файлы из указанной папки в базу.""" p = Path(directory) if not p.is_dir(): print(f"Directory {directory} does not exist.") return for file_path in p.rglob("*.txt") | p.rglob("*.md"): try: content = file_path.read_text(encoding="utf-8") title = file_path.stem add_to_knowledge_base(content, title) print(f"Loaded {file_path}") except Exception as e: print(f"Failed to load {file_path}: {e}") # --------------------------------------------------------------------------- # Интерактивный клиент # --------------------------------------------------------------------------- def interactive_client(): print("Welcome to the RAG agent. Type /help for commands.") while True: try: user_input = input("\n> ") except (EOFError, KeyboardInterrupt): print("\nExiting.") break if not user_input: continue if user_input.startswith("/"): # Команды if user_input.startswith("/add "): path = user_input[5:].strip() if not path: print("Usage: /add ") continue try: content = Path(path).read_text(encoding="utf-8") title = Path(path).stem result = add_to_knowledge_base(content, title) print(result) except Exception as e: print(f"Error reading file: {e}") elif user_input.startswith("/search "): query = user_input[8:].strip() if not query: print("Usage: /search ") continue result = search_knowledge_base(query) print(result) elif user_input in {"/quit", "/exit"}: print("Goodbye!") break elif user_input == "/help": print(textwrap.dedent(""" Commands: /add – add a document to the knowledge base /search – search the knowledge base /quit /exit – exit the program /help – show this help Any other input is treated as a normal user query for the agent. """)) else: print("Unknown command. Type /help for list of commands.") else: # Передаём запрос агенту try: response = agent_executor.invoke({"input": user_input}) print(response["output"]) except Exception as e: print(f"Agent error: {e}") # --------------------------------------------------------------------------- # Точка входа # --------------------------------------------------------------------------- if __name__ == "__main__": # Если передан аргумент – загрузить документы из указанной папки if len(sys.argv) > 1: load_documents_from_dir(sys.argv[1]) interactive_client()