Spaces:
Sleeping
Sleeping
File size: 5,344 Bytes
fd485d9 1e5d06f fd485d9 1e5d06f fd485d9 308de05 1e5d06f 308de05 fd485d9 1e5d06f fd485d9 4a37130 fd485d9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import logging
import os
import re
from typing import List, Optional, Tuple
from pydantic import BaseModel
from common.configuration import Configuration
from components.llm.common import ChatRequest, LlmParams, LlmPredictParams, Message
from components.llm.deepinfra_api import DeepInfraApi
from components.llm.prompts import PROMPT_QE
from components.services.dataset import DatasetService
from components.services.entity import EntityService
from components.services.llm_config import LLMConfigService
logger = logging.getLogger(__name__)
class QEResult(BaseModel):
use_search: bool
search_query: str | None
debug_message: Optional[str | None] = ""
class DialogueService:
def __init__(
self,
config: Configuration,
entity_service: EntityService,
dataset_service: DatasetService,
llm_api: DeepInfraApi,
llm_config_service: LLMConfigService,
) -> None:
self.prompt = PROMPT_QE
self.entity_service = entity_service
self.dataset_service = dataset_service
self.llm_api = llm_api
p = llm_config_service.get_default()
self.llm_params = LlmPredictParams(
temperature=p.temperature,
top_p=p.top_p,
min_p=p.min_p,
seed=p.seed,
frequency_penalty=p.frequency_penalty,
presence_penalty=p.presence_penalty,
n_predict=p.n_predict,
)
async def get_qe_result(self, history: List[Message]) -> QEResult:
"""
Получает результат QE.
Args:
history: История диалога в виде списка сообщений
Returns:
QEResult: Результат QE
"""
request = self._get_qe_request(history)
response = await self.llm_api.predict_chat_stream(
request,
"",
self.llm_params,
)
logger.info(f"QE response: {response}")
try:
return self._postprocess_qe(response)
except Exception as e:
logger.error(f"Error in _postprocess_qe: {e}")
from_chat = self._get_search_query(history)
return QEResult(
use_search=from_chat is not None,
search_query=from_chat.content if from_chat else None,
debug_message=response
)
def get_qe_result_from_chat(self, history: List[Message]) -> QEResult:
from_chat = self._get_search_query(history)
return QEResult(
use_search=from_chat is not None,
search_query=from_chat.content if from_chat else None,
)
def _get_qe_request(self, history: List[Message]) -> ChatRequest:
"""
Подготавливает полный промпт для QE запроса.
Args:
history: История диалога в виде списка сообщений
Returns:
str: Отформатированный промпт с историей диалога
"""
formatted_history = "\n".join(
[self._format_message(msg) for msg in history]
).strip()
message = self.prompt.format(history=formatted_history)
return ChatRequest(
history=[Message(role="user", content=message, searchResults='')]
)
def _format_message(self, message: Message) -> str:
"""
Форматирует сообщение для запроса QE.
Args:
message: Сообщение для форматирования
"""
if message.searchResults:
return f'{message.role}: {message.content}\n<search-results>\n{message.searchResults}\n</search-results>'
return f'{message.role}: {message.content}'
@staticmethod
def _postprocess_qe(input_text: str) -> QEResult:
# Находим все вхождения квадратных скобок
matches = re.findall(r'\[([^\]]*)\]', input_text)
# Проверяем количество найденных скобок
if len(matches) != 2:
raise ValueError("В тексте должно быть ровно две пары квадратных скобок.")
# Извлекаем значения из скобок
first_part = matches[0].strip().lower()
second_part = matches[1].strip()
if first_part == "да":
bool_var = True
elif first_part == "нет":
bool_var = False
else:
raise ValueError("Первая часть текста должна содержать 'ДА' или 'НЕТ'.")
return QEResult(use_search=bool_var, search_query=second_part,
debug_message=input_text)
def _get_search_query(self, history: List[Message]) -> Message | None:
"""
Получает запрос для поиска на основе последнего сообщения пользователя.
"""
return next(
(
msg
for msg in reversed(history)
if msg.role == "user"
and (msg.searchResults is None or not msg.searchResults)
),
None,
)
|