@@ -278,6 +278,9 @@ SPECTACULAR_SETTINGS = { 'REDOC_DIST': 'SIDECAR', } +# Redis +REDIS_HOST=env.str('REDIS_HOST', 'cache-mdb') +REDIS_PORT=env.int('REDIS_PORT', 6379) # Celery CELERY_BROKER_URL = env.str('CELERY_BROKER_URL', 'redis://celery-mdb:6379/0') @@ -3,8 +3,12 @@ import itertools import logging import subprocess import time + +import numpy as np import openpyxl import fitz +import redis +from django.db.models import QuerySet from django.utils.translation import gettext_lazy as _ from datetime import timedelta @@ -32,6 +36,10 @@ from langchain_core.runnables import RunnableWithMessageHistory from langchain_openai.chat_models import ChatOpenAI from langchain_text_splitters import RecursiveCharacterTextSplitter from PIL import Image, UnidentifiedImageError +from redis.commands.search.document import Document +from redis.commands.search.field import TextField, VectorField, TagField +from redis.commands.search.indexDefinition import IndexDefinition, IndexType +from redis.commands.search.query import Query from backend import settings from messages.models import BaseStore, Message @@ -39,9 +47,10 @@ from ml_model.constants import TEMPORARY_TEST_TEXT from ml_model.exceptions import GenerationException from ml_model.models import ( ModelConfiguration, - NeuronModel, + NeuronModel ) from ml_model.services.base import SimpleService +from ml_model.tasks import drop_redis_vectors from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector from poller.models import Proxy @@ -89,6 +98,12 @@ class Chatgpt(SimpleService): }, } + TOOLS_TOKEN_COSTS = { + 'text-embedding-3-large': { + 'output': Decimal('0.000065') + } + } + def __init__(self, store: BaseStore) -> None: super().__init__(store) self.logger = logging.getLogger(self.__class__.__name__) @@ -111,6 +126,7 @@ class Chatgpt(SimpleService): image = None image_size = None normalized_image = None + embedding_tokens = 0 if file: file_extension = Path(file.name).suffix if file_extension == '.pdf': @@ -169,14 +185,21 @@ class Chatgpt(SimpleService): get_session_history=lambda _: chat_history, ) llm_input = [SystemMessage(content=user_system_prompt), HumanMessage(content=input_content)] + input_embedding_tokens = 0 if file and not image: - input_tokens = self.count_text_tokens([*chat_history.messages, *llm_input, *chunks]) + if sum([len(chunk.content) for chunk in chunks]) > 20_000: + input_tokens = self.count_text_tokens([*chat_history.messages, *llm_input, *chunks[:10]]) + input_embedding_tokens = len(chunks) * 600 + else: + input_tokens = self.count_text_tokens([*chat_history.messages, *llm_input, *chunks]) elif image: input_tokens = self.count_text_tokens(llm_input) else: input_tokens = self.count_text_tokens([*chat_history.messages, *llm_input]) output_tokens = 0 - self.assert_enough_balance(input_tokens, image_size, model=self.llm.model_name) + self.assert_enough_balance( + input_tokens, image_size, model=self.llm.model_name, embedding_tokens=input_embedding_tokens + ) if model_name in ('o3-mini', 'gpt-4.5-preview'): system = chat_history.messages.pop(0) messages = [ @@ -228,37 +251,56 @@ class Chatgpt(SimpleService): response = self.llm.invoke(llm_input) chat_history.add_ai_message(response) elif file: - human_messages = [] - chunk_responses = ['Содержание файла: '] - if len(chunks) > 1: - for chunk in chunks: - prompt = [ - { - 'type': 'text', - 'text': f'Сгенерируй 2-3 предложения, которые суммирует следующий текст. ' - f'Включи основную мысль текста и все значимые числовые данные. Текст: {chunk}', - } - ] - human_message = HumanMessage(content=prompt) - human_messages.append(human_message) - response = self.llm.invoke([human_message]) - chunk_responses.append(response.content) + input_tokens = self.count_text_tokens([*chat_history.messages]) + if sum([len(chunk.content) for chunk in chunks]) > 20_000: + redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) + document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + message_uid = str(self.store.messages.first().pk).replace('-', '_') + for chunk_id, chunk in enumerate(chunks): + embedding, e_total_tokens = self.get_embedding(proxy=proxy, content=chunk.content) + self.save_embeddings( + redis_client=redis_client, + message_uid=message_uid, + chunk_id=chunk_id, + text=chunk.content, + embeddings=embedding + ) + embedding_tokens += e_total_tokens + query_embedding, e_total_tokens = self.get_embedding(proxy=proxy, content=input_message.content) + embedding_tokens += e_total_tokens + result = [ + s['section_text'] + for s in self.search_via_embeddings( + redis_client=redis_client, + message_uid=message_uid, + user_query_embeddings=query_embedding + ) + ] + user_input = [ + SystemMessage(content=user_system_prompt), + HumanMessage(self.make_embeddings_prompt( + document_name=document_name, section_texts=result, question=input_message.content + )) + ] + input_tokens += self.count_text_tokens(user_input) + response = conversation.invoke( + {'input': user_input}, + config={'configurable': {'session_id': 'default'}}, + ) + drop_redis_vectors.delay(message_uid) else: - chunk_responses.append(chunks[0].content) - combined_summary = ' '.join(chunk_responses) - question_content = ( - f'Вот краткое содержание каждого чанка:\n' - f'{combined_summary}\n{input_message.content}' - ) - input = [ - SystemMessage(content=user_system_prompt), - HumanMessage(content=f'Используй системный промпт. {question_content}') - ] - input_tokens = self.count_text_tokens(human_messages + input) - response = conversation.invoke( - {'input': input}, - config={'configurable': {'session_id': 'default'}}, - ) + input = [ + SystemMessage(content=user_system_prompt), + HumanMessage( + content=f'Используй системный промпт. Содержание файла: ' + f'{chunks}. Вопрос: {input_message.content}' + ) + ] + input_tokens += self.count_text_tokens(input) + response = conversation.invoke( + {'input': input}, + config={'configurable': {'session_id': 'default'}}, + ) else: # Somehow this chain doesn't support Vision, even though ChatOpenAI (above) does. if model_name in ('o1-preview', 'o1-mini'): @@ -271,10 +313,6 @@ class Chatgpt(SimpleService): if output_tokens == 0: output_tokens = self.count_text_tokens([response]) - if file and not image: - output_tokens += self.count_text_tokens( - [AIMessage(chunk_response) for chunk_response in chunk_responses] - ) if image and normalized_image and model_name not in ('o3-mini', 'gpt-4.5-preview'): self.logger.info(f'Input количество токенов БЕЗ картинки {model_name} - {input_tokens}') @@ -282,7 +320,8 @@ class Chatgpt(SimpleService): self.logger.info(f'Input количество токенов для {model_name} - {input_tokens}') self.logger.info(f'Output количество токенов для {model_name} - {output_tokens}') - self.logger.info(f'Общее количество токенов для {model_name} - {input_tokens + output_tokens}') + self.logger.info(f'Embedding количество токенов для {model_name} - {embedding_tokens}') + self.logger.info(f'Общее количество токенов для {model_name} - {input_tokens + output_tokens + embedding_tokens}') process_time = timedelta(seconds=time.time() - start_time) self.handle_invoice( @@ -290,7 +329,8 @@ class Chatgpt(SimpleService): input_tokens, output_tokens, self.llm.model_name, - info + info, + embedding_tokens ) msgs = self.save_results([response], process_time, save) return msgs @@ -353,14 +393,15 @@ class Chatgpt(SimpleService): input_tokens: int, image_size: tuple | None, model: str = 'gpt-3.5-turbo', + embedding_tokens: int = 0 ): balance = PaymentPlanSelector(self.store.user).get_current_balance() total_tokens = input_tokens - if image_size: total_tokens += self.count_image_tokens(image_size) - input_cost = self.TOKENS_COST[model]['input'] * total_tokens + if embedding_tokens > 0: + input_cost += embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] if input_cost > balance: raise InsufficientBalance(balance, input_cost) @@ -370,6 +411,7 @@ class Chatgpt(SimpleService): output_tokens: int, model: str, info: dict, + embedding_tokens: int = 0, *args, **kwargs, ) -> Decimal: @@ -379,6 +421,8 @@ class Chatgpt(SimpleService): ) if info.get('web_search', 'Отключено') != 'Отключено': price += self.TOKENS_COST[model]['web_search'].get(info.get('web_search', 'medium')) + if embedding_tokens > 0: + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] * embedding_tokens return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def count_image_tokens(self, image_size: tuple, model_version: str = 'gpt-4o') -> int: @@ -491,7 +535,7 @@ class Chatgpt(SimpleService): return 'Файл пуст или содержит изображения, из которых невозможно извлечь текст.' def split_text_to_chunks( - self, raw_text: str, chunk_size: int = 80_000, overlap: int = 300 + self, raw_text: str, chunk_size: int = 4000, overlap: int = 200 ) -> list[HumanMessage]: """ Splitting file raw text to chunks @@ -501,12 +545,89 @@ class Chatgpt(SimpleService): :return: list of chunks """ text_splitter = RecursiveCharacterTextSplitter( - chunk_size=chunk_size, chunk_overlap=overlap, length_function=len + chunk_size=chunk_size, chunk_overlap=overlap, length_function=len, separators=["\n\n", "\n", ".", " ", ""] ) chunks = text_splitter.split_text(raw_text) return [HumanMessage(chunk) for chunk in chunks] - def call_openai_api(self, proxy: Proxy, endpoint: str, json_data: Dict[str, Any]) -> Tuple[Any, Any, AIMessage]: + def get_embedding(self, proxy: Proxy, content: str) -> Tuple[List[float], int]: + ''' + A method for converting raw text (content) into embeddings + using OpenAI API request + :param proxy: Proxy settings object with protocol and address + :param content: raw text of a chunk + ''' + return self.call_openai_api( + proxy=proxy, + endpoint='embeddings', + json_data={ + 'model': 'text-embedding-3-large', + 'input': content + } + ) + + def save_embeddings( + self, redis_client: redis.Redis, message_uid: str, chunk_id: int, text: str, embeddings: List[float] + ) -> None: + ''' + A method for saving embeddings in Redis + :param redis_client: Redis client + :param message_uid: UID of user's message + :param chunk_id: a sequence number of a chunk + :param text: a chunk content + :param embeddings: a list of embeddings getting from a chunk + ''' + embeddings_bytes = np.array(embeddings).astype(dtype=np.float32).tobytes() + redis_client.hset( + f'ml_model:messages:{message_uid}:vectors:{chunk_id}', + mapping={ + 'message_uid': message_uid, + 'section_text': text, + 'section_embeddings': embeddings_bytes + } + ) + + def search_via_embeddings( + self, redis_client: redis.Redis, message_uid: str, user_query_embeddings: List[float], top_k: int = 10 + ) -> List[Document]: + ''' + A method for searching similar vectors to user's query + :param redis_client: Redis client + :param message_uid: UID of user's message + :param user_query_embeddings: a list of embeddings getting from user's query + :param top_k: a number of max return documents + ''' + base_query = f'@message_uid:{{{message_uid}}}=>[KNN {top_k} @section_embeddings $vector AS vector_score]' + query = ( + Query(base_query) + .return_fields('section_text') + .sort_by("vector_score") + .paging(0, top_k) + .dialect(2) + ) + params_dict = {"vector": np.array(user_query_embeddings).astype(dtype=np.float32).tobytes()} + results = redis_client.ft('ml_model-index').search(query, params_dict) + return results.docs + + def make_embeddings_prompt(self, document_name: str, section_texts: List[str], question: str) -> str: + ''' + A method for making a prompt using found embeddings + :param document_name: name of the loaded document + :param section_texts: list of sections' contents + :param question: user question + ''' + return f"""Ты — аналитик данных. Отвечай только на основе предоставленного контекста. + Название файла: {document_name} + Фрагменты: + { + '\n'.join(section_texts) + } + Вопрос: {question} + """ + + def call_openai_api( + self, proxy: Proxy, endpoint: str, json_data: Dict[str, Any] + ) -> Tuple[Any, Any, AIMessage] | Tuple[List[float], int]: ''' A method for sending a request to official openai API :param proxy: Proxy settings object with protocol and address. @@ -540,7 +661,8 @@ class Chatgpt(SimpleService): response = AIMessage(content=content) return input_tokens, output_tokens, response elif ( - (data := resp.json()) + endpoint == 'responses' + and (data := resp.json()) and data.get('output') and ( content := data['output'][-1]['content'][0]['text'] @@ -550,6 +672,13 @@ class Chatgpt(SimpleService): output_tokens = resp.json()['usage']['output_tokens'] response = AIMessage(content=content) return input_tokens, output_tokens, response + elif ( + endpoint == 'embeddings' + and (data := resp.json()) + and (embedding := data['data'][0]['embedding']) + and (tokens := data['usage']['total_tokens']) + ): + return (embedding, tokens) else: raise Exception('GPT not answer correctly, please retry later') @@ -10,7 +10,9 @@ class MLModelConfig(AppConfig): def ready(self): from .signals import create_settings + from .utils import create_redis_search_index setting_changed.connect(create_settings) + create_redis_search_index() return super().ready() @@ -9,6 +9,7 @@ from typing import IO, Any, Dict import deepl import httpx +import redis import replicate import requests from celery import shared_task @@ -192,3 +193,10 @@ def claude_run(payload: dict[str, Any]): @shared_task def evaluate_model(model_name: str, data: Dict[str, Any]): ... + + +@shared_task +def drop_redis_vectors(message_uid: str) -> None: + redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) + for key in redis_client.scan_iter(f'ml_model:messages:{message_uid}:vectors:*'): + redis_client.delete(key) \ No newline at end of file @@ -1,8 +1,13 @@ +import redis import tiktoken from random import randint from typing import Literal, List, Dict, Any, Tuple +from django.conf import settings +from redis.commands.search.field import TagField, TextField, VectorField +from redis.commands.search.indexDefinition import IndexDefinition, IndexType + from authentication.models import CustomUserModel from authentication.selectors.account_status_selector import ( AccountStatusSelector, @@ -55,3 +60,31 @@ def count_openrouter_tokens(model_name: str, messages: List[Dict[str, Any]], out input_tokens += len(encoding.encode(message['content'])) output_tokens = len(encoding.encode(output)) return (input_tokens, output_tokens) + + +def create_redis_search_index() -> None: + ''' + A method for creating an index for storing a chunk's data (content, vectors, etc.) + ''' + redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) + try: + redis_client.ft('ml_model-index').info() + except: + message_uid = TagField('message_uid') + chunk_id = TextField('chunk_id') + section_text = TextField('section_text') + section_embeddings = VectorField( + 'section_embeddings', + 'FLAT', + { + 'TYPE': 'FLOAT32', + 'DIM': 3072, + 'DISTANCE_METRIC': 'COSINE', + 'INITIAL_CAP': 10_000 + } + ) + fields = [message_uid, chunk_id, section_text, section_embeddings] + redis_client.ft('ml_model-index').create_index( + fields=fields, + definition=IndexDefinition(prefix=['ml_model:messages:'], index_type=IndexType.HASH) + ) @@ -3692,7 +3692,6 @@ files = [ {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:bb89f0a835bcfc1d42ccd5f41f04870c1b936d8507c6df12b7737febc40f0909"}, {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:f0c2d907a1e102526dd2986df638343388b94c33860ff3bbe1384130828714b1"}, {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f8157bed2f51db683f31306aa497311b560f2265998122abe1dce6428bd86567"}, - {file = "psycopg2_binary-2.9.10-cp313-cp313-win_amd64.whl", hash = "sha256:27422aa5f11fbcd9b18da48373eb67081243662f9b46e6fd07c3eb46e4535142"}, {file = "psycopg2_binary-2.9.10-cp38-cp38-macosx_12_0_x86_64.whl", hash = "sha256:eb09aa7f9cecb45027683bb55aebaaf45a0df8bf6de68801a6afdc7947bb09d4"}, {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b73d6d7f0ccdad7bc43e6d34273f70d587ef62f824d7261c4ae9b8b1b6af90e8"}, {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ce5ab4bf46a211a8e924d307c1b1fcda82368586a19d0a24f8ae166f5c784864"},