update

fastapi_app.py

@@ -6,7 +6,8 @@ import os
 import datetime
 import json
 import traceback
-from llm.vllm_api import LlmParams
+from llm.common import LlmParams, LlmPredictParams
+from llm.deepinfra_api import DeepInfraApi
 
 # Set the path for log files
 LOGS_BASE_PATH = os.getenv("LOGS_BASE_PATH", "logs")
@@ -17,7 +18,9 @@ LOGS_BASE_PATH = os.getenv("LOGS_BASE_PATH", "logs")
 
 # Check if logs are enabled
 ENABLE_LOGS = os.getenv("ENABLE_LOGS", "0") == "1"
-
+LLM_API_URL = os.getenv("LLM_API_URL", "")
+LLM_API_KEY = os.getenv("LLM_API_KEY", "")
+LLM_USE_DEEPINFRA = os.getenv("LLM_USE_DEEPINFRA", "") == "1"
 
 class Query(BaseModel):
     query: str = ''
@@ -87,18 +90,26 @@ async def search_route(query: Query) -> dict:
 
         llm_params = getattr(query, "llm_params", None)
         
+        if llm_params is None:
+            llm_params = LlmParams(url=LLM_API_URL,api_key=LLM_API_KEY, model="mistralai/Mixtral-8x7B-Instruct-v0.1", predict_params=LlmPredictParams(temperature=0.15, top_p=0.95, min_p=0.05, seed=42, repetition_penalty=1.2, presence_penalty=1.1, max_tokens=6000))
+            
+        if LLM_USE_DEEPINFRA:
+            llm_api = DeepInfraApi(llm_params)
+                                   
+        
         if find_transaction_maps_by_question or find_transaction_maps_by_operation:
-            transaction_maps_results, answer = transaction_maps_search.search_transaction_map(
+            transaction_maps_results, answer = await transaction_maps_search.search_transaction_map(
                 query=question,
                 find_transaction_maps_by_question=find_transaction_maps_by_question,
-                k_neighbours=top)
+                k_neighbours=top,
+                llm_api=llm_api)
 
             response = {'transaction_maps_results': transaction_maps_results}
 
         else:
             modified_query, titles, concat_docs, \
             relevant_consultations, predicted_explanation, \
-            llm_responses = await search.search(question, use_qe, use_olympic, categories, query.llm_params)
+            llm_responses = await search.search(question, use_qe, use_olympic, categories, llm_params)
 
             results = [{'title': str(item1), 'text_for_llm': str(item2)} for item1, item2 in
                         zip(titles, concat_docs)]

llm/common.py

@@ -0,0 +1,64 @@
+from pydantic import BaseModel, Field
+from typing import Optional, List, Protocol
+from abc import ABC, abstractmethod
+
+class LlmPredictParams(BaseModel):
+    """
+    Параметры для предсказания LLM.
+    """
+    system_prompt: Optional[str] = Field(None, description="Системный промпт.")
+    user_prompt: Optional[str] = Field(None, description="Шаблон промпта для передачи от роли user.")
+    n_predict: Optional[int] = None
+    temperature: Optional[float] = None
+    top_k: Optional[int] = None
+    top_p: Optional[float] = None
+    min_p: Optional[float] = None
+    seed: Optional[int] = None
+    repeat_penalty: Optional[float] = None
+    repeat_last_n: Optional[int] = None
+    retry_if_text_not_present: Optional[str] = None
+    retry_count: Optional[int] = None
+    presence_penalty: Optional[float] = None
+    frequency_penalty: Optional[float] = None
+    n_keep: Optional[int] = None
+    cache_prompt: Optional[bool] = None
+    stop: Optional[List[str]] = None
+
+
+class LlmParams(BaseModel):
+    """
+    Основные параметры для LLM.
+    """
+    url: str
+    model: Optional[str] = Field(None, description="Предполагается, что для локального API этот параметр не будет указываться, т.к. будем брать первую модель из списка потому, что модель доступна всего одна. Для deepinfra такой подход не подойдет и модель нужно задавать явно.")
+    type: Optional[str] = None
+    default: Optional[bool] = None
+    template: Optional[str] = None
+    predict_params: Optional[LlmPredictParams] = None
+    api_key: Optional[str] = None
+    
+class LlmApiProtocol(Protocol):
+    async def tokenize(self, prompt: str) -> Optional[dict]:
+        ...
+    async def detokenize(self, tokens: List[int]) -> Optional[str]:
+        ...
+    async def trim_sources(self, sources: str, user_request: str, system_prompt: str = None) -> dict:
+        ...
+    async def predict(self, prompt: str) -> str:
+        ...
+        
+class LlmApi:
+    """
+    Базовый клас для работы с API LLM.
+    """
+    params: LlmParams = None
+    
+    
+    def create_headers(self) -> dict[str, str]:
+        headers = {"Content-Type": "application/json"}
+                
+        if self.params.api_key is not None:
+            headers["Authorization"] = self.params.api_key
+            
+        return headers
+        

llm/deepinfra_api.py

@@ -0,0 +1,153 @@
+import json
+from typing import Optional, List
+import httpx
+from common import LlmPredictParams, LlmParams, LlmApi
+
+class DeepInfraApi(LlmApi):
+    """
+    Класс для работы с API vllm.
+    """
+
+    def __init__(self, params: LlmParams):
+        super.params = params
+
+    
+    async def get_models(self) -> List[str]:
+        """
+        Выполняет GET-запрос к API для получения списка доступных моделей.
+
+        Возвращает:
+            list[str]: Список идентификаторов моделей.
+                       Если произошла ошибка или данные недоступны, возвращается пустой список.
+
+        Исключения:
+            Все ошибки HTTP-запросов логируются в консоль, но не выбрасываются дальше.
+        """
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(f"{super.params.url}/v1/openai/models", super.create_headers())
+                if response.status_code == 200:
+                    json_data = response.json()
+                    return [item['id'] for item in json_data.get('data', [])]
+        except httpx.RequestError as error:
+            print('Error fetching models:', error)
+        return []
+
+    def create_messages(self, prompt: str) -> List[dict]:
+        """
+        Создает сообщения для LLM на основе переданного промпта и системного промпта (если он задан).
+
+        Args:
+            prompt (str): Пользовательский промпт.
+
+        Returns:
+            list[dict]: Список сообщений с ролями и содержимым.
+        """
+        actual_prompt = self.apply_llm_template_to_prompt(prompt)
+        messages = []
+        if super.params.predict_params and super.params.predict_params.system_prompt:
+            messages.append({"role": "system", "content": super.params.predict_params.system_prompt})
+        messages.append({"role": "user", "content": actual_prompt})
+        return messages
+
+    def apply_llm_template_to_prompt(self, prompt: str) -> str:
+        """
+        Применяет шаблон LLM к переданному промпту, если он задан.
+
+        Args:
+            prompt (str): Пользовательский промпт.
+
+        Returns:
+            str: Промпт с примененным шаблоном (или оригинальный, если шаблон отсутствует).
+        """
+        actual_prompt = prompt
+        if super.params.template is not None:
+            actual_prompt = super.params.template.replace("{{PROMPT}}", actual_prompt)
+        return actual_prompt
+
+    async def tokenize(self, prompt: str) -> Optional[dict]:
+        raise NotImplementedError("This function is not supported.")
+
+    async def detokenize(self, tokens: List[int]) -> Optional[str]:
+        raise NotImplementedError("This function is not supported.")
+
+    async def create_request(self, prompt: str) -> dict:
+        """
+        Создает запрос для предсказания на основе параметров LLM.
+
+        Args:
+            prompt (str): Промпт для запроса.
+
+        Returns:
+            dict: Словарь с параметрами для выполнения запроса.
+        """
+
+        request = {
+            "stream": False,
+            "model": super.params.model,
+        }
+
+        predict_params = super.params.predict_params
+        if predict_params:
+            if predict_params.stop:
+                non_empty_stop = list(filter(lambda o: o != "", predict_params.stop))
+                if non_empty_stop:
+                    request["stop"] = non_empty_stop
+                    
+            if predict_params.n_predict is not None:
+                request["max_tokens"] = int(predict_params.n_predict or 0)
+                
+            request["temperature"] = float(predict_params.temperature or 0)
+            if predict_params.top_k is not None:
+                request["top_k"] = int(predict_params.top_k)
+                
+            if predict_params.top_p is not None:
+                request["top_p"] = float(predict_params.top_p)
+                
+            if predict_params.min_p is not None:
+                request["min_p"] = float(predict_params.min_p)
+                
+            if predict_params.seed is not None:
+                request["seed"] = int(predict_params.seed)
+                
+            if predict_params.n_keep is not None:
+                request["n_keep"] = int(predict_params.n_keep)
+                
+            if predict_params.cache_prompt is not None:
+                request["cache_prompt"] = bool(predict_params.cache_prompt)
+                
+            if predict_params.repeat_penalty is not None:
+                request["repetition_penalty"] = float(predict_params.repeat_penalty)
+                
+            if predict_params.repeat_last_n is not None:
+                request["repeat_last_n"] = int(predict_params.repeat_last_n)
+                
+            if predict_params.presence_penalty is not None:
+                request["presence_penalty"] = float(predict_params.presence_penalty)
+                
+            if predict_params.frequency_penalty is not None:
+                request["frequency_penalty"] = float(predict_params.frequency_penalty)
+
+        request["messages"] = self.create_messages(prompt)
+        return request
+
+    async def trim_sources(self, sources: str, user_request: str, system_prompt: str = None) -> dict:
+        raise NotImplementedError("This function is not supported.")
+        
+    async def predict(self, prompt: str) -> str:
+        """
+        Выполняет запрос к API и возвращает результат.
+
+        Args:
+            prompt (str): Входной текст для предсказания.
+
+        Returns:
+            str: Сгенерированный текст.
+        """
+        async with httpx.AsyncClient() as client:
+            request = await self.create_request(prompt)
+
+            async with httpx.AsyncClient() as client:
+                response = client.post(f"{super.params.url}/v1/openai/chat/completions", super.create_headers(), json=request)
+                if response.status_code == 200:
+                    return response.json()["choices"][0]["message"]["content"]

llm/vllm_api.py

@@ -3,51 +3,17 @@ from typing import Optional, List, Any
 
 import httpx
 from pydantic import BaseModel, Field
+from common import LlmPredictParams, LlmParams, LlmApi
 
 
-class LlmPredictParams(BaseModel):
-    """
-    Параметры для предсказания LLM.
-    """
-    system_prompt: Optional[str] = Field(None, description="Системный промпт.")
-    user_prompt: Optional[str] = Field(None, description="Шаблон промпта для передачи от роли user.")
-    n_predict: Optional[int] = None
-    temperature: Optional[float] = None
-    top_k: Optional[int] = None
-    top_p: Optional[float] = None
-    min_p: Optional[float] = None
-    seed: Optional[int] = None
-    repeat_penalty: Optional[float] = None
-    repeat_last_n: Optional[int] = None
-    retry_if_text_not_present: Optional[str] = None
-    retry_count: Optional[int] = None
-    presence_penalty: Optional[float] = None
-    frequency_penalty: Optional[float] = None
-    n_keep: Optional[int] = None
-    cache_prompt: Optional[bool] = None
-    stop: Optional[List[str]] = None
-
-
-class LlmParams(BaseModel):
-    """
-    Основные параметры для LLM.
-    """
-    url: str
-    type: Optional[str] = None
-    default: Optional[bool] = None
-    template: Optional[str] = None
-    predict_params: Optional[LlmPredictParams] = None
-
-
-class LlmApi:
+class LlmApi(LlmApi):
     """
     Класс для работы с API vllm.
     """
-    params: LlmParams = None
 
     def __init__(self, params: LlmParams):
-        self.params = params
-
+        super.params = params
+    
     async def get_models(self) -> List[str]:
         """
         Выполняет GET-запрос к API для получения списка доступных моделей.
@@ -61,13 +27,26 @@ class LlmApi:
         """
         try:
             async with httpx.AsyncClient() as client:
-                response = await client.get(f"{self.params.url}/v1/models", headers={"Content-Type": "application/json"})
+                response = await client.get(f"{super.params.url}/v1/models", super.create_headers())
                 if response.status_code == 200:
                     json_data = response.json()
                     return [item['id'] for item in json_data.get('data', [])]
         except httpx.RequestError as error:
             print('Error fetching models:', error)
         return []
+    
+    async def get_model(self) -> str:
+        model = None
+        if super.params.model is not None:
+            model = super.params.model
+        else:
+            models = await self.get_models()
+            model = models[0] if models else None
+            
+        if model is None:
+            raise Exception("No model name provided and no models available.")
+        
+        return model
 
     def create_messages(self, prompt: str) -> List[dict]:
         """
@@ -81,8 +60,8 @@ class LlmApi:
         """
         actual_prompt = self.apply_llm_template_to_prompt(prompt)
         messages = []
-        if self.params.predict_params and self.params.predict_params.system_prompt:
-            messages.append({"role": "system", "content": self.params.predict_params.system_prompt})
+        if super.params.predict_params and super.params.predict_params.system_prompt:
+            messages.append({"role": "system", "content": super.params.predict_params.system_prompt})
         messages.append({"role": "user", "content": actual_prompt})
         return messages
 
@@ -97,8 +76,8 @@ class LlmApi:
             str: Промпт с примененным шаблоном (или оригинальный, если шаблон отсутствует).
         """
         actual_prompt = prompt
-        if self.params.template is not None:
-            actual_prompt = self.params.template.replace("{{PROMPT}}", actual_prompt)
+        if super.params.template is not None:
+            actual_prompt = super.params.template.replace("{{PROMPT}}", actual_prompt)
         return actual_prompt
 
     async def tokenize(self, prompt: str) -> Optional[dict]:
@@ -112,14 +91,10 @@ class LlmApi:
             Optional[dict]: Словарь с токенами и максимальной длиной модели, если запрос успешен.
                             Если запрос неуспешен, возвращает None.
         """
-        model = (await self.get_models())[0] if await self.get_models() else None
-        if not model:
-            print("No models available for tokenization.")
-            return None
 
         actual_prompt = self.apply_llm_template_to_prompt(prompt)
         request_data = {
-            "model": model,
+            "model": self.get_model(),
             "prompt": actual_prompt,
             "add_special_tokens": False,
         }
@@ -127,9 +102,9 @@ class LlmApi:
         try:
             async with httpx.AsyncClient() as client:
                 response = await client.post(
-                    f"{self.params.url}/tokenize",
+                    f"{super.params.url}/tokenize",
                     json=request_data,
-                    headers={"Content-Type": "application/json"},
+                    headers=super.create_headers(),
                 )
                 if response.status_code == 200:
                     data = response.json()
@@ -155,19 +130,15 @@ class LlmApi:
             Optional[str]: Строка, полученная в результате детокенизации, если запрос успешен.
                            Если запрос неуспешен, возвращает None.
         """
-        model = (await self.get_models())[0] if await self.get_models() else None
-        if not model:
-            print("No models available for detokenization.")
-            return None
-
-        request_data = {"model": model, "tokens": tokens or []}
+        
+        request_data = {"model": self.get_model(), "tokens": tokens or []}
 
         try:
             async with httpx.AsyncClient() as client:
                 response = await client.post(
-                    f"{self.params.url}/detokenize",
+                    f"{super.params.url}/detokenize",
                     json=request_data,
-                    headers={"Content-Type": "application/json"},
+                    headers=super.create_headers(),
                 )
                 if response.status_code == 200:
                     data = response.json()
@@ -192,17 +163,14 @@ class LlmApi:
         Returns:
             dict: Словарь с параметрами для выполнения запроса.
         """
-        models = await self.get_models()
-        if not models:
-            raise ValueError("No models available to create a request.")
-        model = models[0]
+        model = self.get_model()
 
         request = {
             "stream": True,
             "model": model,
         }
 
-        predict_params = self.params.predict_params
+        predict_params = super.params.predict_params
         if predict_params:
             if predict_params.stop:
                 non_empty_stop = list(filter(lambda o: o != "", predict_params.stop))
@@ -283,7 +251,7 @@ class LlmApi:
         # Максимально допустимое количество токенов для источников
         max_length = (
             max_token_count
-            - (self.params.predict_params.n_predict or 0)
+            - (super.params.predict_params.n_predict or 0)
             - aux_token_count
             - system_prompt_token_count
         )
@@ -322,7 +290,7 @@ class LlmApi:
             request = await self.create_request(prompt)
 
             # Начинаем потоковый запрос
-            async with client.stream("POST", f"{self.params.url}/v1/chat/completions", json=request) as response:
+            async with client.stream("POST", f"{super.params.url}/v1/chat/completions", json=request) as response:
                 if response.status_code != 200:
                     # Если ошибка, читаем ответ для получения подробностей
                     error_content = await response.aread()

transaction_maps_search.py

@@ -3,14 +3,13 @@ from business_transaction_map.common.constants import DEVICE, DO_NORMALIZATION,
 from business_transaction_map.components.faiss_vector_database import FaissVectorDatabase
 from business_transaction_map.components.embedding_extraction import EmbeddingExtractor
 import os
-import requests
 from prompts import BUSINESS_TRANSACTION_PROMPT
+from llm.common import LlmApi
 
 
 db_files_path = os.environ.get("GLOBAL_TRANSACTION_MAPS_DATA_PATH", "transaction_maps_search_data/csv/карта_проводок_new.pkl")
 
 model_path = os.environ.get("GLOBAL_TRANSACTION_MAPS_MODEL_PATH", "")
-llm_api_endpoint = os.environ.get("LLM_API_ENDPOINT", "")
 
 class TransactionMapsSearch:
     
@@ -26,14 +25,11 @@ class TransactionMapsSearch:
         self.database = FaissVectorDatabase(str(db_files_path))
         
     @staticmethod
-    def extract_business_transaction_with_llm(question: str) -> str:
+    async def extract_business_transaction_with_llm(question: str, llm_api: LlmApi) -> str:
+        prompt = BUSINESS_TRANSACTION_PROMPT.replace('{{ЗАПРОС}}', question)
+        res = await llm_api.predict(prompt)
         
-        question = BUSINESS_TRANSACTION_PROMPT.replace('{{ЗАПРОС}}', question)
-
-        response = requests.post(url=llm_api_endpoint,
-                                 json={"prompt": f"[INST] {question} [/INST]", #пробелы внутри [INST], как оказалось, важны. Без них можно словить бесконечную генерацию бреда от ллм
-                                       "temperature": 0.0})
-        return response.json()['content']
+        return res
         
 
     @staticmethod
@@ -66,13 +62,14 @@ class TransactionMapsSearch:
         return answer
 
 
-    def search_transaction_map(self, 
+    async def search_transaction_map(self, 
                                query: str = None,
                                find_transaction_maps_by_question: bool = False,
-                               k_neighbours: int = 15):
+                               k_neighbours: int = 15,
+                               llm_api: LlmApi = None):
         
         if find_transaction_maps_by_question:
-            query = self.extract_business_transaction_with_llm(query)
+            query = await self.extract_business_transaction_with_llm(query, llm_api)
         cleaned_text = query.replace("\n", " ")
         # cleaned_text = 'query: ' + cleaned_text  # only for e5
         query_tokens = self.model.query_tokenization(cleaned_text)