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

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG-памятью/agent.py

@@ -0,0 +1,185 @@
+import os
+import json
+import logging
+from typing import List, Dict, Any
+
+from langchain_ollama import OllamaEmbeddings, ChatOllama
+from langchain_qdrant import QdrantStore
+from langchain_core.tools import tool
+from langchain.agents import AgentExecutor, ZeroShotAgent, Tool
+from langchain.schema import HumanMessage, AIMessage
+
+# ==========================
+# Конфигурация и логирование
+# ==========================
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
+QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "knowledge_base")
+
+# ==========================
+# Векторный хранилище
+# ==========================
+
+try:
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+except Exception as e:
+    logger.error(f"Не удалось инициализировать OllamaEmbeddings: {e}")
+    raise
+
+try:
+    qdrant_store = QdrantStore(
+        url=QDRANT_URL,
+        collection_name=QDRANT_COLLECTION,
+        embeddings=embeddings,
+    )
+except Exception as e:
+    logger.error(f"Не удалось подключиться к Qdrant: {e}")
+    raise
+
+
+# ==========================
+# Инструменты
+# ==========================
+
+@tool("search_knowledge_base")
+def search_knowledge_base(query: str) -> List[Dict[str, Any]]:
+    """
+    Поиск релевантных документов в базе знаний.
+
+    Args:
+        query (str): поисковый запрос.
+
+    Returns:
+        list of dict: список найденных документов с полями 'id', 'text' и 'score'.
+    """
+    try:
+        results = qdrant_store.similarity_search_with_score(query, k=5)
+    except Exception as e:
+        logger.exception(f"Ошибка при поиске в Qdrant: {e}")
+        return []
+
+    docs = []
+    for doc, score in results:
+        docs.append(
+            {
+                "id": getattr(doc.metadata, "doc_id", None),
+                "text": doc.page_content,
+                "score": score,
+            }
+        )
+    logger.debug(f"search_knowledge_base: найдено {len(docs)} документов")
+    return docs
+
+
+@tool("add_to_knowledge_base")
+def add_to_knowledge_base(text: str) -> str:
+    """
+    Добавление нового текста в базу знаний.
+
+    Args:
+        text (str): текст, который нужно добавить.
+
+    Returns:
+        str: сообщение об успехе или ошибке.
+    """
+    try:
+        # Разбиваем на чанки с гибкой схемой
+        chunk_size = 500  # символов
+        overlap = 50      # символов перекрытия
+
+        chunks = []
+        start = 0
+        while start < len(text):
+            end = min(start + chunk_size, len(text))
+            chunk_text = text[start:end]
+            chunks.append(chunk_text)
+            start += chunk_size - overlap
+
+        for i, chunk in enumerate(chunks):
+            doc_id = f"{hash(chunk)}_{i}"
+            qdrant_store.add_documents(
+                documents=[chunk],
+                ids=[doc_id],
+                metadata={"doc_id": doc_id},
+            )
+    except Exception as e:
+        logger.exception(f"Ошибка при добавлении в Qdrant: {e}")
+        return "Не удалось добавить документ."
+
+    logger.info(f"add_to_knowledge_base: добавлено {len(chunks)} чанков")
+    return f"Документ успешно добавлен, {len(chunks)} чанков."
+
+
+# ==========================
+# Агент
+# ==========================
+
+tools = [
+    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="Добавьте новый текст в базу знаний.",
+    ),
+]
+
+prompt = ZeroShotAgent.create_prompt(
+    tools,
+    prefix="""
+Вы — интеллектуальный агент, помогающий пользователю находить и добавлять информацию в базу знаний.
+Используйте инструменты search_knowledge_base и add_to_knowledge_base по мере необходимости.
+""",
+)
+
+llm = ChatOllama(model="llama3")
+
+agent_chain = ZeroShotAgent(llm=llm, tools=tools, prompt=prompt)
+agent_executor = AgentExecutor.from_agent_and_tools(
+    agent=agent_chain,
+    tools=tools,
+    verbose=True,
+)
+
+
+# ==========================
+# Тесты (можно запускать как `python agent.py`)
+# ==========================
+
+def _test_add_and_search():
+    test_text = (
+        "Python — это высокоуровневый язык программирования, "
+        "который поддерживает несколько парадигм и широко используется в научных исследованиях."
+    )
+    add_result = add_to_knowledge_base(test_text)
+    assert "успешно" in add_result.lower(), f"Добавление не удалось: {add_result}"
+
+    search_query = "Python язык программирования"
+    results = search_knowledge_base(search_query)
+    assert isinstance(results, list), "Результат поиска должен быть списком."
+    assert len(results) > 0, "Ни один документ не найден."
+
+    logger.info("Тесты пройдены успешно.")
+
+
+if __name__ == "__main__":
+    # Пример использования агента
+    user_query = input("Введите ваш запрос: ")
+    try:
+        response = agent_executor.invoke({"input": user_query})
+        print("\nОтвет агента:")
+        print(response["output"])
+    except Exception as e:
+        logger.exception(f"Ошибка при выполнении запроса: {e}")
+
+    # Запуск тестов
+    _test_add_and_search()