136 lines
4.7 KiB
Python
136 lines
4.7 KiB
Python
import os
|
||
import logging
|
||
from typing import Any, Dict, List
|
||
|
||
from langchain_ollama import OllamaEmbeddings, ChatOllama
|
||
from langchain_qdrant import QdrantStore
|
||
from langchain_core.documents import Document
|
||
from langchain_core.runnables import RunnableLambda
|
||
|
||
# Настройка логирования
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Конфигурация подключения к Qdrant
|
||
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
|
||
COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "knowledge_base")
|
||
|
||
# Инициализация эмбеддингов и модели Ollama
|
||
try:
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
except Exception as e:
|
||
logger.exception("Не удалось подключиться к Ollama Embeddings: %s", e)
|
||
raise
|
||
|
||
try:
|
||
llm = ChatOllama(model="llama3")
|
||
except Exception as e:
|
||
logger.exception("Не удалось подключиться к Ollama LLM: %s", e)
|
||
raise
|
||
|
||
# Инициализация хранилища Qdrant
|
||
try:
|
||
qdrant_store = QdrantStore(
|
||
url=QDRANT_URL,
|
||
collection_name=COLLECTION_NAME,
|
||
embeddings=embeddings,
|
||
)
|
||
except Exception as e:
|
||
logger.exception("Не удалось подключиться к Qdrant: %s", e)
|
||
raise
|
||
|
||
|
||
# ==========================
|
||
# Middleware для обработки ошибок и логирования
|
||
# ==========================
|
||
def middleware(func):
|
||
"""
|
||
Декоратор, который оборачивает функцию в try/except,
|
||
логирует входные параметры и результат.
|
||
"""
|
||
|
||
def wrapper(*args, **kwargs):
|
||
func_name = func.__name__
|
||
logger.debug("Вызов %s с args=%s kwargs=%s", func_name, args, kwargs)
|
||
try:
|
||
result = func(*args, **kwargs)
|
||
logger.debug("Результат %s: %s", func_name, result)
|
||
return result
|
||
except Exception as exc:
|
||
logger.exception("Ошибка в %s: %s", func_name, exc)
|
||
raise
|
||
|
||
return wrapper
|
||
|
||
|
||
# ==========================
|
||
# RAG инструменты
|
||
# ==========================
|
||
@middleware
|
||
def search_knowledge_base(query: str) -> List[Document]:
|
||
"""
|
||
Поиск релевантных документов в Qdrant.
|
||
Возвращает список Document с полями content и metadata.
|
||
"""
|
||
if not query:
|
||
raise ValueError("Запрос не может быть пустым")
|
||
|
||
try:
|
||
docs = qdrant_store.similarity_search(query, k=5)
|
||
logger.info("Найдено %d документов по запросу '%s'", len(docs), query)
|
||
return docs
|
||
except Exception as exc:
|
||
logger.exception("Ошибка при поиске: %s", exc)
|
||
raise
|
||
|
||
|
||
@middleware
|
||
def add_to_knowledge_base(content: str, metadata: Dict[str, Any] | None = None) -> None:
|
||
"""
|
||
Добавляет новый документ в Qdrant.
|
||
content – текст документа,
|
||
metadata – произвольные метаданные (может быть None).
|
||
"""
|
||
if not content:
|
||
raise ValueError("Содержимое не может быть пустым")
|
||
|
||
try:
|
||
doc = Document(page_content=content, metadata=metadata or {})
|
||
qdrant_store.add_documents([doc])
|
||
logger.info("Документ добавлен в коллекцию '%s'", COLLECTION_NAME)
|
||
except Exception as exc:
|
||
logger.exception("Ошибка при добавлении документа: %s", exc)
|
||
raise
|
||
|
||
|
||
# ==========================
|
||
# Пример использования (можно использовать как тесты)
|
||
# ==========================
|
||
if __name__ == "__main__":
|
||
# Минимальный набор unit‑тестов
|
||
import unittest
|
||
|
||
class TestAgent(unittest.TestCase):
|
||
def test_add_and_search(self):
|
||
test_content = "Тестовый документ для проверки RAG."
|
||
test_metadata = {"source": "unit_test"}
|
||
|
||
# Добавляем документ
|
||
add_to_knowledge_base(test_content, test_metadata)
|
||
|
||
# Ищем по ключевому слову
|
||
results = search_knowledge_base("тестовый")
|
||
self.assertTrue(any(test_content in doc.page_content for doc in results))
|
||
|
||
def test_empty_query(self):
|
||
with self.assertRaises(ValueError):
|
||
search_knowledge_base("")
|
||
|
||
def test_empty_content(self):
|
||
with self.assertRaises(ValueError):
|
||
add_to_knowledge_base("")
|
||
|
||
unittest.main(argv=["first-arg-is-ignored"], exit=False) |