165 lines
5.2 KiB
Python
165 lines
5.2 KiB
Python
import sys
|
|
import logging
|
|
from typing import List
|
|
|
|
# LangChain imports
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from langchain.schema 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
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
LLM_MODEL = "llama3"
|
|
|
|
# Создаём объект для получения эмбеддингов
|
|
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
|
|
|
# Подключаемся к локальному экземпляру Qdrant (по умолчанию localhost:6333)
|
|
vector_store = QdrantVectorStore(
|
|
embedding_function=embeddings,
|
|
url="http://localhost:6333",
|
|
collection_name="documents",
|
|
)
|
|
|
|
|
|
def middleware(func):
|
|
"""
|
|
Декоратор для обработки ошибок и логирования.
|
|
"""
|
|
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
logger.info(f"Выполняется команда: {func.__name__}")
|
|
result = func(*args, **kwargs)
|
|
logger.info("Команда выполнена успешно")
|
|
return result
|
|
except Exception as e:
|
|
logger.exception(f"Ошибка в команде {func.__name__}: {e}")
|
|
print(f"[ERROR] {e}")
|
|
|
|
return wrapper
|
|
|
|
|
|
@middleware
|
|
def add_document(text: str) -> None:
|
|
"""
|
|
Добавляет документ в Qdrant.
|
|
"""
|
|
doc = Document(page_content=text)
|
|
vector_store.add_documents([doc])
|
|
print("Документ добавлен.")
|
|
|
|
|
|
@middleware
|
|
def search(query: str, k: int = 3) -> List[Document]:
|
|
"""
|
|
Выполняет поиск по запросу и выводит результаты.
|
|
"""
|
|
results = vector_store.similarity_search_with_score(query, k=k)
|
|
if not results:
|
|
print("Ничего не найдено.")
|
|
return []
|
|
|
|
for idx, (doc, score) in enumerate(results, start=1):
|
|
print(f"\nРезультат {idx} (score: {score:.4f})")
|
|
print(doc.page_content[:500] + ("..." if len(doc.page_content) > 500 else ""))
|
|
return [doc for doc, _ in results]
|
|
|
|
|
|
def parse_command(line: str):
|
|
"""
|
|
Разбирает строку команды и вызывает соответствующую функцию.
|
|
"""
|
|
if not line.strip():
|
|
return
|
|
|
|
parts = line.split(maxsplit=1)
|
|
cmd = parts[0].lower()
|
|
|
|
if cmd == "/add":
|
|
if len(parts) < 2:
|
|
print("[ERROR] Необходимо указать текст для добавления.")
|
|
return
|
|
add_document(parts[1])
|
|
elif cmd == "/search":
|
|
if len(parts) < 2:
|
|
print("[ERROR] Необходимо указать запрос для поиска.")
|
|
return
|
|
search(parts[1])
|
|
elif cmd == "/quit":
|
|
print("Выход из программы.")
|
|
sys.exit(0)
|
|
else:
|
|
print(f"[ERROR] Нераспознанная команда: {cmd}")
|
|
|
|
|
|
def main():
|
|
"""
|
|
Основной цикл CLI.
|
|
"""
|
|
logger.info("Запуск клиента Human-in-the-Loop")
|
|
print("Добро пожаловать! Используйте команды /add, /search и /quit.")
|
|
while True:
|
|
try:
|
|
line = input("> ")
|
|
parse_command(line)
|
|
except KeyboardInterrupt:
|
|
print("\nПрервано пользователем. Выход.")
|
|
break
|
|
except EOFError:
|
|
print("\nEOF. Выход.")
|
|
break
|
|
|
|
|
|
# -------------------- UNIT TESTS --------------------
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
class TestClient(unittest.TestCase):
|
|
def setUp(self):
|
|
# Подменяем vector_store для тестов
|
|
self.original_vector_store = vector_store
|
|
mock_vs = MagicMock()
|
|
mock_vs.add_documents.return_value = None
|
|
mock_vs.similarity_search_with_score.return_value = [
|
|
(Document(page_content="Test content 1"), 0.9),
|
|
(Document(page_content="Test content 2"), 0.8),
|
|
]
|
|
global vector_store
|
|
vector_store = mock_vs
|
|
|
|
def tearDown(self):
|
|
global vector_store
|
|
vector_store = self.original_vector_store
|
|
|
|
@patch('builtins.print')
|
|
def test_add_document_success(self, mock_print):
|
|
add_document("Sample text")
|
|
mock_print.assert_called_with("Документ добавлен.")
|
|
|
|
@patch('builtins.print')
|
|
def test_search_results(self, mock_print):
|
|
results = search("query")
|
|
self.assertEqual(len(results), 2)
|
|
# Проверяем, что print был вызван хотя бы один раз
|
|
self.assertTrue(mock_print.call_count > 0)
|
|
|
|
@patch('builtins.print')
|
|
def test_parse_unknown_command(self, mock_print):
|
|
parse_command("/unknown")
|
|
mock_print.assert_called_with("[ERROR] Нераспознанная команда: /unknown")
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) > 1 and sys.argv[1] == "test":
|
|
# Запуск тестов
|
|
unittest.main(argv=[sys.argv[0]])
|
|
else:
|
|
main() |