diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/rag_agent.py new file mode 100644 index 0000000..0e09ed4 --- /dev/null +++ b/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 , /search , /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() \ No newline at end of file