diff --git a/main.py b/main.py index a4ab016..4af42f0 100644 --- a/main.py +++ b/main.py @@ -1,16 +1,23 @@ # main.py # Полностью рабочий пример агента с RAG‑памятью на базе Qdrant и Ollama. -# Используется LangChain 1.x, create_agent, инструменты @tool, и Ollama‑LLM/embeddings. +# Весь код написан в одном файле для простоты демонстрации. +# +# Требования: +# - Python 3.10+ +# - Qdrant (работает по умолчанию на localhost:6333) +# - Ollama (llama3 + nomic-embed-text) +# - LangChain 1.2.10+ и связанные пакеты # # Запуск: -# python main.py -# После запуска можно использовать команды: -# /add <content> – добавить документ -# /search <query> <max> – семантический поиск -# /quit – выйти +# 1. Убедитесь, что Qdrant и Ollama запущены. +# 2. pip install -r requirements.txt +# 3. python main.py # -# Для загрузки документов из директории используйте функцию load_documents_from_dir. -#""" +# После запуска появится интерактивный клиент с командами: +# /add <path_to_file> – добавить документ в базу знаний. +# /search <query> – выполнить поиск. +# /quit – выйти. +# Любой другой ввод – агент обработает как обычный запрос. import os import sys @@ -22,14 +29,14 @@ 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, AgentToolkit, Tool +from langchain.agents import create_agent, AgentExecutor, AgentType from langchain_core.messages import HumanMessage # --------------------------------------------------------------------------- # Конфигурация # --------------------------------------------------------------------------- -QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") -QDRANT_COLLECTION = "knowledge_base" +QDRANT_URL = "http://localhost:6333" +COLLECTION_NAME = "knowledge" EMBEDDING_MODEL = "nomic-embed-text" LLM_MODEL = "llama3" @@ -40,141 +47,152 @@ LLM_MODEL = "llama3" embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) vector_store = QdrantVectorStore( url=QDRANT_URL, - collection_name=QDRANT_COLLECTION, + collection_name=COLLECTION_NAME, embedding=embeddings, ) # --------------------------------------------------------------------------- # Чанкинг # --------------------------------------------------------------------------- -text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) +text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) # --------------------------------------------------------------------------- # Инструменты # --------------------------------------------------------------------------- -@tool("search_knowledge_base", "Semantic search in the knowledge base.") +@tool("search_knowledge_base", "Semantic search in the knowledge base") def search_knowledge_base(query: str, max_results: int = 5) -> str: - """Return top‑k relevant documents for a query. - The function returns a formatted string with titles and snippets. - """ + """Возвращает топ‑N релевантных фрагментов из Qdrant.""" results = vector_store.similarity_search_with_score(query, k=max_results) if not results: return "No relevant documents found." + # Формируем читаемый ответ formatted = [] - for doc, score in results: - title = doc.metadata.get("title", "Untitled") - snippet = doc.page_content[:200].replace("\n", " ") - formatted.append(f"{title} (score: {score:.3f}): {snippet}...") + 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: - """Chunk the content, embed, and store in Qdrant. - Returns a confirmation message. - """ +@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) - docs = [] - for i, chunk in enumerate(chunks): - docs.append( - { - "page_content": chunk, - "metadata": {"title": title, "chunk_index": i}, - } - ) + # Создаём 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(chunks)} chunks of '{title}' to the knowledge base." + 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) -# Список инструментов -tools = [search_knowledge_base, add_to_knowledge_base] - # Создаём агент agent = create_agent( llm=llm, - tools=tools, - system_message="You are an assistant that can search and add documents to a local knowledge base. Use the provided tools.", - verbose=True, + 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=tools, verbose=True) +# Обёртка для удобного вызова +agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True) # --------------------------------------------------------------------------- -# Загрузка документов из директории +# Загрузка документов из директории (инициализация) # --------------------------------------------------------------------------- def load_documents_from_dir(directory: str) -> None: - """Load all .txt files from a directory into the knowledge base. - Each file becomes a separate document with its filename as title. - """ - path = Path(directory) - if not path.is_dir(): + """Загружает все .txt и .md файлы из указанной папки в базу.""" + p = Path(directory) + if not p.is_dir(): print(f"Directory {directory} does not exist.") return - for file in path.glob("*.txt"): - title = file.stem - content = file.read_text(encoding="utf-8") - print(f"Adding {title}...", end=" ") - result = add_to_knowledge_base(content, title) - print(result) + 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}") # --------------------------------------------------------------------------- -# CLI +# Интерактивный клиент # --------------------------------------------------------------------------- -def main(): - # Если пользователь передал путь к директории, загрузим документы - if len(sys.argv) > 1: - load_documents_from_dir(sys.argv[1]) - - print("\n--- RAG Agent CLI ---") - print("Commands:") - print(" /add <title> <content> – add a document") - print(" /search <query> <max> – search knowledge base") - print(" /quit – exit") - +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.strip(): + 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 <path_to_file>") + 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 <query>") + 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 <path> – add a document to the knowledge base + /search <q> – 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 user_input.startswith("/quit"): - print("Goodbye!") - break - - if user_input.startswith("/add"): - parts = user_input.split(maxsplit=2) - if len(parts) < 3: - print("Usage: /add <title> <content>") - continue - title, content = parts[1], parts[2] - print(add_to_knowledge_base(content, title)) - continue - - if user_input.startswith("/search"): - parts = user_input.split(maxsplit=2) - if len(parts) < 2: - print("Usage: /search <query> [max_results]") - continue - query = parts[1] - max_results = int(parts[2]) if len(parts) > 2 else 5 - print(search_knowledge_base(query, max_results)) - continue - - # Любой другой ввод – передаём агенту - response = agent_executor.invoke({"input": user_input}) - print(response.get("output", "")) - - +# --------------------------------------------------------------------------- +# Точка входа +# --------------------------------------------------------------------------- if __name__ == "__main__": - main() + # Если передан аргумент – загрузить документы из указанной папки + if len(sys.argv) > 1: + load_documents_from_dir(sys.argv[1]) + interactive_client()