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

Всем привет! Это Игорь Густомясов, 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.

После серии экспериментов мы остановились на на 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.