From a03f0812cd18a343c17b203cb677d40215fcc580 Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Thu, 18 Jun 2026 13:09:53 +0000 Subject: [PATCH] main.py written --- main.py | 277 ++++++++++++++++++++++++++------------------------------ 1 file changed, 126 insertions(+), 151 deletions(-) diff --git a/main.py b/main.py index 4af42f0..2dc9a82 100644 --- a/main.py +++ b/main.py @@ -1,27 +1,11 @@ -# 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 – выйти. -# Любой другой ввод – агент обработает как обычный запрос. +"""Main module implementing a RAG-enabled agent with Qdrant and Ollama. + +The code follows the assignment specification and uses only the required libraries. +""" import os import sys -import textwrap +import json from pathlib import Path from typing import List, Dict, Any @@ -29,170 +13,161 @@ 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 +from langchain.agents import create_agent, AgentExecutor +from langchain.schema import AgentAction, AgentFinish +from langchain_core.prompts import ChatPromptTemplate # --------------------------------------------------------------------------- -# Конфигурация +# Configuration # --------------------------------------------------------------------------- -QDRANT_URL = "http://localhost:6333" -COLLECTION_NAME = "knowledge" +QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") +QDRANT_COLLECTION = "knowledge" EMBEDDING_MODEL = "nomic-embed-text" LLM_MODEL = "llama3" +CHUNK_SIZE = 1000 # characters +CHUNK_OVERLAP = 200 # --------------------------------------------------------------------------- -# Векторное хранилище +# Vector store helper # --------------------------------------------------------------------------- -# Инициализируем эмбеддер и клиент Qdrant -embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) -vector_store = QdrantVectorStore( - url=QDRANT_URL, - collection_name=COLLECTION_NAME, - embedding=embeddings, -) +class KnowledgeBase: + """Wrapper around QdrantVectorStore providing add/search helpers.""" + + def __init__(self, url: str = QDRANT_URL, collection: str = QDRANT_COLLECTION): + self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) + self.store = QdrantVectorStore.from_existing_collection( + collection_name=collection, + url=url, + embedding=self.embeddings, + ) + + def add_documents(self, documents: List[str], titles: List[str]): + """Adds a list of documents with corresponding titles to the store. + + Each document is split into chunks before being stored. + """ + splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP) + for doc, title in zip(documents, titles): + chunks = splitter.split_text(doc) + metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))] + self.store.add_texts(chunks, metadatas=metadatas) + + def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]: + """Semantic search in the knowledge base. + + Returns a list of dicts with keys: text, title, distance. + """ + results = self.store.similarity_search_with_score(query, k=max_results) + return [ + { + "text": text, + "title": meta.get("title", "unknown"), + "distance": score, + } + for text, score, meta in results + ] # --------------------------------------------------------------------------- -# Чанкинг +# Global knowledge base instance # --------------------------------------------------------------------------- -text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) +kb = KnowledgeBase() # --------------------------------------------------------------------------- -# Инструменты +# Tools # --------------------------------------------------------------------------- -@tool("search_knowledge_base", "Semantic search in the knowledge base") +@tool("search_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) + """Search the knowledge base for a query and return formatted results.""" + results = kb.search(query, 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) + formatted = [f"Title: {r['title']}\nSnippet: {r['text'][:200]}...\nDistance: {r['distance']:.4f}" for r in results] + return "\n\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}'." +@tool("add_to_knowledge_base") +def add_to_knowledge_base(content: str, title: str) -> str: + """Add a new document to the knowledge base.""" + kb.add_documents([content], [title]) + return f"Document '{title}' added successfully." # --------------------------------------------------------------------------- -# Агент +# Agent definition # --------------------------------------------------------------------------- -# Системный промпт, который подсказывает агенту использовать инструменты -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, +SYSTEM_PROMPT = ( + "You are an assistant that can search and add documents to a local knowledge base. " + "Use the provided tools to perform semantic search and store new information. " + "When answering user queries, first determine if the user needs a search or an addition. " + "If no relevant information is found, suggest adding new content." ) -# Обёртка для удобного вызова -agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True) +prompt = ChatPromptTemplate.from_messages([ + ("system", SYSTEM_PROMPT), + ("user", "{input}"), +]) + +agent = create_agent( + llm=ChatOllama(model=LLM_MODEL), + tools=[search_knowledge_base, add_to_knowledge_base], + prompt=prompt, + verbose=True, +) + +agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base]) # --------------------------------------------------------------------------- -# Загрузка документов из директории (инициализация) +# CLI client # --------------------------------------------------------------------------- +def load_documents_from_dir(directory: str): + """Load all .txt files from a directory and add them to the knowledge base.""" + docs = [] + titles = [] + for path in Path(directory).glob("**/*.txt"): + text = path.read_text(encoding="utf-8") + docs.append(text) + titles.append(path.stem) + if docs: + kb.add_documents(docs, titles) + print(f"Loaded {len(docs)} documents from {directory}.") + else: + print("No .txt files found.") -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.") +def interactive_loop(): + print("RAG Agent CLI. Commands: /add <content>, /search <query>, /load <dir>, /quit") while True: try: - user_input = input("\n> ") - except (EOFError, KeyboardInterrupt): + user_input = input("> ") + except (KeyboardInterrupt, EOFError): 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 <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: - # Передаём запрос агенту + if user_input.startswith("/quit"): + print("Goodbye!") + break + elif user_input.startswith("/add "): try: - response = agent_executor.invoke({"input": user_input}) - print(response["output"]) - except Exception as e: - print(f"Agent error: {e}") + _, rest = user_input.split("/add ", 1) + title, content = rest.split(" ", 1) + response = add_to_knowledge_base(content, title) + print(response) + except ValueError: + print("Usage: /add <title> <content>") + elif user_input.startswith("/search "): + query = user_input.split("/search ", 1)[1] + results = search_knowledge_base(query) + print(results) + elif user_input.startswith("/load "): + dir_path = user_input.split("/load ", 1)[1] + load_documents_from_dir(dir_path) + else: + # Treat as normal user query + result = agent_executor.invoke({"input": user_input}) + print(result.get("output", "")) -# --------------------------------------------------------------------------- -# Точка входа -# --------------------------------------------------------------------------- if __name__ == "__main__": - # Если передан аргумент – загрузить документы из указанной папки - if len(sys.argv) > 1: - load_documents_from_dir(sys.argv[1]) - interactive_client() + if len(sys.argv) > 1 and sys.argv[1] == "cli": + interactive_loop() + else: + print("Usage: python main.py cli") + print("Run the interactive CLI with: python main.py cli")