Умный поиск по API, или NLP против функционального поиска

f19ebe0dba68f661751070bec6de4bd2.png

Всем привет! Это Игорь Густомясов, CTO кластера техноплатформы в МТС, и Никита Бояндин, ведущий разработчик в том же кластере. (Да, мы создали текст вместе.) Рассказываем о поиске данных API для Интеграционной платформы МТС.

Наш коллега Александр Бардаш круто расписал, как мы развиваем функции Интеграционной платформы. Так вот: получилось настолько хорошо, что возникла проблема.

В экосистеме МТС множество продуктов — от проката самокатов до высокотехнологичных сервисов The Platform. Стоило интеграционной платформе встать на ноги, как на ней резко выросло количество спецификаций API.

Так перед нами развернулась двойная задача: не только технически поддержать различные протоколы взаимодействия (HTTPS, gRPC, GraphQL и прочие), но и сделать поиск данных API. Решение — под катом.

Гипотезы и первые попытки

Как правило, в подобных случаях используют каталогизацию и поиск по ней. Но нам этот подход не помог.

Дело в том, что при переходе определенного порога по количеству API каталогизация становится непрозрачной для клиентов. Даже если ее выстроить каноничным, архитектурно правильным путем — пользователь далеко не всегда знает глобальный контекст своей локальной задачи. Его цель проста: найти нужный интерфейс и начать его использовать. Чем быстрее получается, тем удобнее Интеграционная платформа, вот и всё.

Ок, родилась вторая гипотеза. Современные спецификации хорошо структурированы и содержат поля для документирования. Почему бы не прикрутить полнотекстовый поиск?

Мы взяли open-source движок и проиндексировали наши спецификации. Столкнулись с другой проблемой. Многие API разрабатывались давно и задокументированы плохо. В этом случае полнотекстовый поиск бессилен, если не вычищать данные. Дело нужное, но тогда нам не хватило бы времени развивать платформу.

Кроме того, язык документирования внутренних API бывает довольно специфичен и привязан к продуктовому или технологическому домену. Порой сложновато сформировать правильный запрос.

В третий раз закинули мы невод: попробовали мощь современных больших языковых моделей (LLM) и смежных технологий. И кажется, поймали золотую рыбку.

Однако искать оптимальное решение было непросто. Пришлось сделать несколько макетов, обкатать на практике, провести серию технологических итераций…, но обо всем по порядку.

Первый подход и RetrivalQA

Задача была в новинку, и мы испробовали множество методов поиска. Направление-то мы сразу взяли верное: библиотека LangChain. Но в API встречаются конфиденциальные данные. Значит, внешняя LLM не подходила, только внутренняя.

Мы написали собственный class, чтобы совместить LLM с LangChain. Оставалось сократить время поиска до секунды или ниже — пользователь не хочет ждать. И мы… не справились, так как любой запрос в LLM по API отвечает не менее чем за 2 секунды. Увы (

Вот наш class для локального использования LLM:

class CustomLLM(LLM):
    def _call(self, prompt, stop=None):
        url = 'your_url'
        headers = {
            "Authorization": f"Bearer dummy", 
            "Content-Type": "application/json"
        }
    
        payload = {
            "model": "LLM_MODEL",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.2
        }

        response = requests.post(url, headers=headers, json=payload)
        if response.status_code == 200:
            return response.json()['choices'][0]['message']['content']
        else:
            logger.error(f"Request failed with status code {response.status_code}")
            return {"error": f"Request failed with status code {response.status_code}"}
    
    @property
    def _identifying_params(self):
        return {"model": "LLM_MODEL"}
    
    @property
    def _llm_type(self):
        return "custom"


custom_llm = CustomLLM()

qa_chain = RetrievalQA.from_chain_type(llm=custom_llm, 
                                  chain_type="stuff", 
                                  retriever=retriever, 
                                  return_source_documents=True)

Токенизация

Еще мы столкнулись с проблемой токенизации данных. В основном поисковые запросы на русском языке. А все API на английском. Какой же моделью делать векторы?

Пробовали русские rubert и bi-encoder. К сожалению, спустя долгие часы валидации и попыток улучшить скор мы констатировали, что эти модели не подошли.

Тогда мы вернулись к самому простому варианту с использованием легкой англоязычной модели:

embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

Выбор базы для данных

Давайте спустимся на уровень данных и посмотрим, где же хранить наши эмбеддинги. Напрашиваются векторные базы: Qdrant, ChromaDB и Weaviate.

Аспект

Qdrant

ChromaDB

Weaviate

Тип решения

Поддерживает метаданные

С открытым исходным кодом

С графовыми данными и метаданными

Поддержка моделей

Без моделей векторизации, но интегрируется с внешними

С моделями для генерации векторов: GPT и Sentence-BERT

Интегрируется с моделями BERT, OpenAI и другими

Поиск

HNSW для быстрого поиска по схожести

HNSW для поиска по схожести

HNSW и другие методы для поиска по схожести и фильтрации

Гибкость

Высокая, но требует настроить внешние модели

Средняя, но просто использовать

Очень высокая, поддерживает сложные запросы и графы

Масштабируемость

Отличная для больших данных

Хорошая

Отличная, поддерживает горизонтальное масштабирование

Сложность

Простая настройка, но требуется интегрировать модели

Низкая, но недостает гибкости

Сложная настройка, ресурсоемкая, требует инфраструктуры

После серии экспериментов мы остановились на на Weaviate, поскольку эта база прекрасно себя показала в решении задача QA по пространствам Confluence. Сейчас в нем, помимо простого текста, мы храним вложения и осуществляем по ним поиск. Также Weaviate дает большие возможности кастомизации поиска, но об этом как-нибудь в другой раз.

Ensemble Retrivers и наше конечное решение

Мы ушли в FAISS (хранение данных в MongoDB): просто использовать, выдает нужные метрики. Впоследствии взяли технику ансамблевого ретривера из FAISS и BM25, что нас и спасло. Вот часть кода:

from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

doc_list_1 = [
    "I like apples",
    "I like oranges",
    "Apples and oranges are fruits",
]

# initialize the bm25 retriever and faiss retriever
bm25_retriever = BM25Retriever.from_texts(
    doc_list_1, metadatas=[{"source": 1}] * len(doc_list_1)
)
bm25_retriever.k = 2

doc_list_2 = [
    "You like apples",
    "You like oranges",
]

embedding = OpenAIEmbeddings()
faiss_vectorstore = FAISS.from_texts(
    doc_list_2, embedding, metadatas=[{"source": 2}] * len(doc_list_2)
)
faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={"k": 2})

# initialize the ensemble retriever
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, faiss_retriever], weights=[0.5, 0.5]
)

Дальнейшие задачи

У любого семантического поиска множество проблем: орфография и исправление слов, разные языки и так далее. А идеальный результат все еще должен выходить не более чем за секунду.

Но мы уверены, что справимся, ведь наш подход перспективен. Мы можем добавлять на Интеграционную платформу новые технологии и решения — и эта возможность сохранится.

Надеемся, материал был полезен и интересен. До скорых встреч на поле LLM)

© Habrahabr.ru