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

2026-05-27 10:24:08 +00:00
parent 8c60256b9d
commit 3f811b0bc2
1 changed files with 143 additions and 0 deletions
@@ -0,0 +1,143 @@
+import os
+import sys
+import logging
+from typing import List
+
+from langchain_ollama import ChatOllama
+from langchain_qdrant import QdrantStore
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+
+# Настройка логирования
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+logger = logging.getLogger(__name__)
+
+# Конфигурация Ollama и Qdrant
+OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
+EMBEDDING_MODEL = os.getenv("NOMIC_EMBED_TEXT_MODEL", "nomic-embed-text")
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
+COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection")
+
+# Инициализация Ollama и Qdrant
+try:
+    llm = ChatOllama(model=OLLAMA_MODEL)
+except Exception as e:
+    logger.error(f"Не удалось подключиться к Ollama: {e}")
+    sys.exit(1)
+
+try:
+    qdrant_store = QdrantStore(
+        url=QDRANT_URL,
+        collection_name=COLLECTION_NAME,
+        embedding_function=lambda text: llm.embed_query(text),
+    )
+except Exception as e:
+    logger.error(f"Не удалось подключиться к Qdrant: {e}")
+    sys.exit(1)
+
+# Текстовый разделитель с гибкой схемой chunking
+def split_text(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Document]:
+    """
+    Делит входной текст на чанки.
+    Параметры:
+        chunk_size: максимальная длина чанка в символах
+        overlap: количество пересекающихся символов между соседними чанками
+    Возвращает список Document объектов.
+    """
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=overlap,
+        separators=["\n\n", "\n", " ", ""],
+    )
+    return [Document(page_content=c) for c in splitter.split_text(text)]
+
+def add_document(file_path: str):
+    """Добавляет документ из файла в базу Qdrant."""
+    if not os.path.isfile(file_path):
+        logger.warning(f"Файл не найден: {file_path}")
+        return
+
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            content = f.read()
+    except Exception as e:
+        logger.error(f"Ошибка чтения файла {file_path}: {e}")
+        return
+
+    docs = split_text(content)
+    if not docs:
+        logger.warning("Документ пустой после разбиения на чанки.")
+        return
+
+    try:
+        qdrant_store.add_documents(docs)
+        logger.info(f"Добавлено {len(docs)} чанков из файла: {file_path}")
+    except Exception as e:
+        logger.error(f"Ошибка при добавлении документов в Qdrant: {e}")
+
+def search(query: str, k: int = 5):
+    """Ищет релевантные документы и выводит ответы."""
+    try:
+        docs_and_scores = qdrant_store.similarity_search_with_score(query, k=k)
+    except Exception as e:
+        logger.error(f"Ошибка при поиске в Qdrant: {e}")
+        return
+
+    if not docs_and_scores:
+        logger.info("Ничего не найдено.")
+        return
+
+    # Формируем контекст из найденных чанков
+    context = "\n\n".join([doc.page_content for doc, _ in docs_and_scores])
+
+    prompt = f"""
+Ты эксперт по теме. На основе следующего контекста ответь на вопрос:
+
+Контекст:
+{context}
+
+Вопрос: {query}
+"""
+
+    try:
+        response = llm.invoke(prompt)
+        print("\nОтвет:\n" + response.content.strip())
+    except Exception as e:
+        logger.error(f"Ошибка при генерации ответа Ollama: {e}")
+
+def main():
+    logger.info("Запущен CLI клиент RAG. Команды: /add <путь>, /search <запрос>, /quit")
+    while True:
+        try:
+            user_input = input("\n> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nВыход.")
+            break
+
+        if not user_input:
+            continue
+
+        if user_input.lower() == "/quit":
+            logger.info("Завершение работы.")
+            break
+        elif user_input.startswith("/add"):
+            parts = user_input.split(maxsplit=1)
+            if len(parts) != 2:
+                print("Использование: /add <путь_к_файлу>")
+                continue
+            add_document(parts[1])
+        elif user_input.startswith("/search"):
+            parts = user_input.split(maxsplit=1)
+            if len(parts) != 2:
+                print("Использование: /search <запрос>")
+                continue
+            search(parts[1])
+        else:
+            print("Неизвестная команда. Доступные команды: /add, /search, /quit")
+
+if __name__ == "__main__":
+    main()