@@ -42,13 +42,12 @@ class FileProcessingService: @classmethod def get_pdf_data(cls, pdf_data: bytes) -> str: try: - doc = fitz.open(stream=pdf_data, filetype='pdf') raw_text = '' - for page_number, page in enumerate(doc, start=1): - content = page.get_text('text') - if content: - raw_text += content - doc.close() + with fitz.open(stream=pdf_data, filetype='pdf') as doc: + for page_number, page in enumerate(doc, start=1): + content = page.get_text('text') + if content: + raw_text += content fitz.TOOLS.store_shrink(100) except Exception: return f'Ошибка: Файл поврежден или не может быть прочитан.' @@ -138,145 +138,175 @@ class Chatgpt(SimpleService): else: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) for proxy in Proxy.objects.all(): - self.llm = ChatOpenAI( - model=model_name, - http_client=httpx.Client(proxy=f'{proxy.protocol}://{proxy.address}'), - ) - self.llm.temperature = info.pop('temperature', 0.5) - self.llm.model_kwargs = { - 'presence_penalty': info.pop('presence', 0), - 'top_p': info.pop('top_p', 0.5), - } - self.llm.tiktoken_model_name = 'gpt-4' - chat_history = self.get_chat_history(model_name=model_name) - chat_history.add_message(HumanMessage(content=input_message.content)) - conversation = RunnableWithMessageHistory( - runnable=self.llm, - get_session_history=lambda _: chat_history, - ) - llm_input = [SystemMessage(content=user_system_prompt), HumanMessage(content=input_content)] - input_tokens, input_embedding_tokens = self._get_input_tokens( - file, image, chunks, chat_history, llm_input, model_name - ) - output_tokens = 0 - self.assert_enough_balance( - input_tokens, image_size, model=self.llm.model_name, embedding_tokens=input_embedding_tokens - ) - if model_name == 'gpt-oss-120b': - system = chat_history.messages.pop(0) - messages = [ - { - 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', - 'content': msg.content, - } - for msg in chat_history.messages - ] - messages.insert(0, {'role': 'system', 'content': system.content}) - messages.insert(0, {'role': 'system', 'content': user_system_prompt}) - if file and not image: - if sum([len(chunk.content) for chunk in chunks]) > 20_000: - document_name = ( - chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - ) - embedding_tokens, file_data = EmbeddingService.get_large_file_data( - self.store.messages.first().pk, - text_chunks, - proxy, - input_message.content, - model='text-embedding-3-small', - index_name='ml_model-index-1536', - ) - messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( - document_name=document_name, - section_texts=file_data, - question=input_message.content, - ) - else: - messages[-1]['content'] = ( - 'Используй системный промпт. Содержание файла: ' - f'{"".join(text_chunks)}. Вопрос: {input_message.content}' - ) - json_data = {'model': f'openai/{model_name}', 'messages': messages} - response = httpx.post( - url='https://openrouter.ai/api/v1/chat/completions', - proxy=f'{proxy.protocol}://{proxy.address}', - headers={'Authorization': f'Bearer {settings.OPENROUTER_API_KEY}'}, - timeout=600, - json=json_data, + http_client = httpx.Client(proxy=f'{proxy.protocol}://{proxy.address}') + try: + self.llm = ChatOpenAI( + model=model_name, + http_client=http_client, ) - if ( - (data := response.json()) - and data.get('choices') - and ( - content := ','.join( - [choice['message']['content'] for choice in data.get('choices')]) - ) - ): - input_tokens = response.json()['usage']['prompt_tokens'] - output_tokens = response.json()['usage']['completion_tokens'] - response = AIMessage(content=content.replace('\\n', '\n')) - else: - raise Exception('GPT not answer correctly, please retry later') - elif model_name == 'o3-mini': - system = chat_history.messages.pop(0) - messages = [ - { - 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', - 'content': msg.content, - } - for msg in chat_history.messages - ] - messages.insert(0, {'role': 'system', 'content': system.content}) - messages.insert(0, {'role': 'system', 'content': user_system_prompt}) - if image: - messages[-1]['content'] = [ - {'type': 'text', 'text': input_message.content}, - image_data, + self.llm.temperature = info.pop('temperature', 0.5) + self.llm.model_kwargs = { + 'presence_penalty': info.pop('presence', 0), + 'top_p': info.pop('top_p', 0.5), + } + self.llm.tiktoken_model_name = 'gpt-4' + chat_history = self.get_chat_history(model_name=model_name) + chat_history.add_message(HumanMessage(content=input_message.content)) + conversation = RunnableWithMessageHistory( + runnable=self.llm, + get_session_history=lambda _: chat_history, + ) + llm_input = [SystemMessage(content=user_system_prompt), HumanMessage(content=input_content)] + input_tokens, input_embedding_tokens = self._get_input_tokens( + file, image, chunks, chat_history, llm_input, model_name + ) + output_tokens = 0 + self.assert_enough_balance( + input_tokens, image_size, model=self.llm.model_name, embedding_tokens=input_embedding_tokens + ) + if model_name == 'gpt-oss-120b': + system = chat_history.messages.pop(0) + messages = [ + { + 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', + 'content': msg.content, + } + for msg in chat_history.messages ] - elif file: - if sum([len(chunk.content) for chunk in chunks]) > 20_000: - document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = EmbeddingService.get_large_file_data( - self.store.messages.first().pk, - text_chunks, - proxy, - input_message.content, - model='text-embedding-3-small', - index_name='ml_model-index-1536', - ) - messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( - document_name=document_name, - section_texts=file_data, - question=input_message.content, + messages.insert(0, {'role': 'system', 'content': system.content}) + messages.insert(0, {'role': 'system', 'content': user_system_prompt}) + if file and not image: + if sum([len(chunk.content) for chunk in chunks]) > 20_000: + document_name = ( + chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + ) + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + json_data = {'model': f'openai/{model_name}', 'messages': messages} + response = httpx.post( + url='https://openrouter.ai/api/v1/chat/completions', + proxy=f'{proxy.protocol}://{proxy.address}', + headers={'Authorization': f'Bearer {settings.OPENROUTER_API_KEY}'}, + timeout=600, + json=json_data, + ) + if ( + (data := response.json()) + and data.get('choices') + and ( + content := ','.join( + [choice['message']['content'] for choice in data.get('choices')]) ) + ): + input_tokens = response.json()['usage']['prompt_tokens'] + output_tokens = response.json()['usage']['completion_tokens'] + response = AIMessage(content=content.replace('\\n', '\n')) else: - messages[-1]['content'] = ( - 'Используй системный промпт. Содержание файла: ' - f'{"".join(text_chunks)}. Вопрос: {input_message.content}' - ) - json_data = { - 'model': model_name, - 'messages': messages - } - input_tokens, output_tokens, response = self.call_openai_api(proxy=proxy, endpoint='chat/completions',json_data=json_data) - elif info.get('web_search', 'Отключено') != 'Отключено': - system = chat_history.messages.pop(0) - messages = [ - {'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', 'content': msg.content} - for msg in chat_history.messages - ] - messages.insert(0, {'role': 'system', 'content': system.content}) - messages.insert(0, {'role': 'system', 'content': user_system_prompt}) - search_context_size, json_data = self.get_web_search_data( - info.get('web_search', 'Средний контекст'), model_name, messages - ) - info['web_search'] = search_context_size - if image: - messages[-1]['content'] = [ - {'type': 'input_text', 'text': input_message.content}, - {'type': 'input_image', 'image_url': image_data['image_url']['url']}, + raise Exception('GPT not answer correctly, please retry later') + elif model_name == 'o3-mini': + system = chat_history.messages.pop(0) + messages = [ + { + 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', + 'content': msg.content, + } + for msg in chat_history.messages + ] + messages.insert(0, {'role': 'system', 'content': system.content}) + messages.insert(0, {'role': 'system', 'content': user_system_prompt}) + if image: + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + image_data, + ] + elif file: + if sum([len(chunk.content) for chunk in chunks]) > 20_000: + document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + json_data = { + 'model': model_name, + 'messages': messages + } + input_tokens, output_tokens, response = self.call_openai_api(proxy=proxy, endpoint='chat/completions',json_data=json_data) + elif info.get('web_search', 'Отключено') != 'Отключено': + system = chat_history.messages.pop(0) + messages = [ + {'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', 'content': msg.content} + for msg in chat_history.messages ] + messages.insert(0, {'role': 'system', 'content': system.content}) + messages.insert(0, {'role': 'system', 'content': user_system_prompt}) + search_context_size, json_data = self.get_web_search_data( + info.get('web_search', 'Средний контекст'), model_name, messages + ) + info['web_search'] = search_context_size + if image: + messages[-1]['content'] = [ + {'type': 'input_text', 'text': input_message.content}, + {'type': 'input_image', 'image_url': image_data['image_url']['url']}, + ] + elif file: + if sum([len(chunk.content) for chunk in chunks]) > 20_000: + document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + input_tokens, output_tokens, response = self.call_openai_api( + proxy=proxy, endpoint='responses', json_data=json_data + ) + elif image: + response = self.llm.invoke(llm_input) + chat_history.add_ai_message(response) elif file: + input_tokens = self.count_text_tokens([*chat_history.messages]) if sum([len(chunk.content) for chunk in chunks]) > 20_000: document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] embedding_tokens, file_data = EmbeddingService.get_large_file_data( @@ -287,98 +317,73 @@ class Chatgpt(SimpleService): model='text-embedding-3-small', index_name='ml_model-index-1536', ) - messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( - document_name=document_name, - section_texts=file_data, - question=input_message.content, + user_input = [ + SystemMessage(content=user_system_prompt), + HumanMessage( + EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + ), + ] + input_tokens += self.count_text_tokens(user_input) + response = conversation.invoke( + {'input': user_input}, + config={'configurable': {'session_id': 'default'}}, ) else: - messages[-1]['content'] = ( - 'Используй системный промпт. Содержание файла: ' - f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + input = [ + SystemMessage(content=user_system_prompt), + HumanMessage( + content=( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + ), + ] + input_tokens += self.count_text_tokens(input) + response = conversation.invoke( + {'input': input}, + config={'configurable': {'session_id': 'default'}}, ) - input_tokens, output_tokens, response = self.call_openai_api( - proxy=proxy, endpoint='responses', json_data=json_data - ) - elif image: - response = self.llm.invoke(llm_input) - chat_history.add_ai_message(response) - elif file: - input_tokens = self.count_text_tokens([*chat_history.messages]) - if sum([len(chunk.content) for chunk in chunks]) > 20_000: - document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = EmbeddingService.get_large_file_data( - self.store.messages.first().pk, - text_chunks, - proxy, - input_message.content, - model='text-embedding-3-small', - index_name='ml_model-index-1536', - ) - user_input = [ - SystemMessage(content=user_system_prompt), - HumanMessage( - EmbeddingService.make_embeddings_prompt( - document_name=document_name, - section_texts=file_data, - question=input_message.content, - ) - ), - ] - input_tokens += self.count_text_tokens(user_input) - response = conversation.invoke( - {'input': user_input}, - config={'configurable': {'session_id': 'default'}}, - ) else: - input = [ - SystemMessage(content=user_system_prompt), - HumanMessage( - content=( - 'Используй системный промпт. Содержание файла: ' - f'{"".join(text_chunks)}. Вопрос: {input_message.content}' - ) - ), - ] - input_tokens += self.count_text_tokens(input) response = conversation.invoke( - {'input': input}, + {'input': llm_input}, config={'configurable': {'session_id': 'default'}}, ) - else: - response = conversation.invoke( - {'input': llm_input}, - config={'configurable': {'session_id': 'default'}}, - ) - chat_history.add_ai_message(response) + chat_history.add_ai_message(response) - if output_tokens == 0: - output_tokens = self.count_text_tokens([response]) + if output_tokens == 0: + output_tokens = self.count_text_tokens([response]) - if ( - image - and normalized_image - and model_name != 'o3-mini' - ): - self.logger.info(f'Input количество токенов БЕЗ картинки {model_name} - {input_tokens}') - input_tokens += self.count_image_tokens(normalized_image.size, model_name) - - self.logger.info(f'Input количество токенов для {model_name} - {input_tokens}') - self.logger.info(f'Output количество токенов для {model_name} - {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( - input_message.content_object.model, - input_tokens, - output_tokens, - self.llm.model_name, - info, - embedding_tokens - ) - msgs = self.save_results([response], process_time, save) - return msgs + if ( + image + and normalized_image + and model_name != 'o3-mini' + ): + self.logger.info(f'Input количество токенов БЕЗ картинки {model_name} - {input_tokens}') + input_tokens += self.count_image_tokens(normalized_image.size, model_name) + + self.logger.info(f'Input количество токенов для {model_name} - {input_tokens}') + self.logger.info(f'Output количество токенов для {model_name} - {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( + input_message.content_object.model, + input_tokens, + output_tokens, + self.llm.model_name, + info, + embedding_tokens + ) + msgs = self.save_results([response], process_time, save) + http_client.close() + return msgs + finally: + http_client.close() def get_chat_history(self, model_name: str) -> InMemoryChatMessageHistory: if isinstance(self.store, Chat): @@ -516,13 +521,14 @@ class Chatgpt(SimpleService): return total_tokens def _get_image_data(self, file_bytes: bytes, file_extension: str) -> Tuple: - normalized_image = Image.open(BytesIO(file_bytes)).convert('RGB') - buf = BytesIO() + with Image.open(BytesIO(file_bytes)) as source_image: + normalized_image = source_image.convert('RGB') format = 'jpeg' if file_extension not in ('png', 'jpeg', 'webp') else file_extension - normalized_image.save(buf, format=format) - image_url = f'data:image/{format};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:image/{format};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' image_size = normalized_image.size + normalized_image.close() image_data = {'type': 'image_url', 'image_url': {'url': image_url}} return normalized_image, image_size, image_data @@ -114,12 +114,11 @@ class Claude(SimpleService): else: kind = filetype.guess(file_bytes[:20]) mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(file) format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(file) as normalized_image: + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' messages[-1]['content'] = [ {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, @@ -89,14 +89,15 @@ class Flux_2(SimpleService): extension = kind.extension if extension.upper() not in (extensions := ['JPG', 'JPEG', 'PNG', 'WEBP']): raise FileExtensionNotSupported(extensions) - normalized_image = Image.open(BytesIO(file_bytes)).convert('RGB') - buf = BytesIO() format = 'jpeg' if extension not in ('png', 'jpeg', 'webp') else extension - normalized_image.save(buf, format=format) - file_width, file_height = get_image_dimensions(buf) - input_mp = math.ceil((file_width*file_height) / 1_000_000) - image = f'data:image/{format};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(BytesIO(file_bytes)) as source_image: + normalized_image = source_image.convert('RGB') + with BytesIO() as buf: + normalized_image.save(buf, format=format) + file_width, file_height = get_image_dimensions(buf) + input_mp = math.ceil((file_width*file_height) / 1_000_000) + image = f'data:image/{format};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + normalized_image.close() callback_data.update({'input_images': [image]}) images = [replicate_run(f'black-forest-labs/{version}', callback_data)] process_time = timedelta(seconds=(time.time() - start_time)) @@ -43,15 +43,14 @@ class Fluxlorafast(SimpleService): messages: list[Message] = [] for image in images: for proxy in Proxy.objects.all(): - client = httpx.Client( - base_url='https://queue.fal.run', - headers={'Authorization': f'Key {settings.FAL_API_KEY}'}, - timeout=600, - proxy=f'{proxy.protocol}://{proxy.address}', - ) - try: - file = File(BytesIO(client.get(image).content), '.png') + with httpx.Client( + base_url='https://queue.fal.run', + headers={'Authorization': f'Key {settings.FAL_API_KEY}'}, + timeout=600, + proxy=f'{proxy.protocol}://{proxy.address}', + ) as client: + file = File(BytesIO(client.get(image).content), '.png') except Exception: continue @@ -111,40 +110,39 @@ class Fluxlorafast(SimpleService): proxies = Proxy.objects.all() for proxy in proxies: requests_number = 0 - client = httpx.Client( + with httpx.Client( base_url='https://queue.fal.run', headers={'Authorization': f'Key {settings.FAL_API_KEY}'}, timeout=600, proxy=f'{proxy.protocol}://{proxy.address}', - ) - result = client.post( - f'fal-ai/{version}', - json={'prompt': input_message.content, **callback_data}, - ).json() - - is_success = False - - while True: - if requests_number == 271 / len(proxies): - break - - try: - status = client.get(result['status_url']).json() - except Exception: - requests_number += 1 - continue - if status.get('status') == 'COMPLETED': - is_success = True + ) as client: + result = client.post( + f'fal-ai/{version}', + json={'prompt': input_message.content, **callback_data}, + ).json() + + is_success = False + + while True: + if requests_number == 271 / len(proxies): + break + + try: + status = client.get(result['status_url']).json() + except Exception: + requests_number += 1 + continue + if status.get('status') == 'COMPLETED': + is_success = True + break + + time.sleep(1 / 3) + + if is_success: + final_result = client.get(result['response_url']).json() break - time.sleep(1 / 3) - - if is_success: - break - if not is_success: raise GenerationException from ModelTimeoutError - final_result = client.get(result['response_url']).json() - return final_result @@ -168,12 +168,11 @@ class Gemini(SimpleService): ) elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(file) format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(file) as normalized_image: + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' messages[-1]['content'] = [ {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, @@ -126,12 +126,11 @@ class Gemini_3_1(SimpleService): elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): kind = filetype.guess(file_bytes[:20]) mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(input_message.file) format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(input_message.file) as normalized_image: + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' messages[-1]['content'] = [ {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, @@ -104,9 +104,11 @@ class Gptimage(SimpleService): raise CorruptedFileError normalized_image = BytesIO(file_bytes) if kind.extension.upper() != 'PNG': - img = Image.open(normalized_image).convert('RGBA') + with Image.open(normalized_image) as source_image: + img = source_image.convert('RGBA') normalized_image = BytesIO() img.save(normalized_image, format='PNG') + img.close() normalized_image.seek(0) files = {'image': ('image.png', normalized_image, 'image/png')} if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( @@ -122,12 +122,11 @@ class Grok(SimpleService): elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): mime = kind.mime if kind else 'application/octet-stream' input_message.file.seek(0) - normalized_image = Image.open(input_message.file) format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(input_message.file) as normalized_image: + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' messages[-1]['content'] = [ {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, @@ -73,12 +73,11 @@ class Llama(SimpleService): if image: kind = filetype.guess(image.read(20)) mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(image) as normalized_image: + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' messages[-1]['content'] = [ {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, @@ -62,12 +62,11 @@ class Mistral(SimpleService): if image: kind = filetype.guess(image.read(20)) mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) format = 'jpeg' if kind and kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(image) as normalized_image: + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' messages[-1]['content'] = [ {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, @@ -104,12 +104,11 @@ class Qwen_3_6(SimpleService): elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): kind = filetype.guess(file_bytes[:20]) mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(input_message.file) format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() + with Image.open(input_message.file) as normalized_image: + with BytesIO() as buf: + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' messages[-1]['content'] = [ {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, @@ -91,31 +91,32 @@ class Raifgpt(Chatgpt): text_chunks = EmbeddingService.split_text_to_chunks(text, chunk_size=1000) chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] for proxy in Proxy.objects.all(): - self.llm = ChatOpenAI( - model='gpt-4o', - http_client=httpx.Client(proxy=f'{proxy.protocol}://{proxy.address}'), - ) - self.llm.tiktoken_model_name = 'gpt-4' - self.llm.temperature = 0.8 - self.llm.top_p = 1 - self.llm.presence_penalty = 0 - chat_history = self.get_chat_history(model_name='gpt-4o') - conversation = RunnableWithMessageHistory( - runnable=self.llm, - get_session_history=lambda _: chat_history, - ) - llm_input = HumanMessage(content=input_content) - if file: - image_count = getattr(self, 'image_count', 0) - input_tokens = self.count_text_tokens([*chat_history.messages, llm_input, *chunks]) - input_tokens += Decimal(image_count) * Decimal('0.13') - else: - input_tokens = self.count_text_tokens([*chat_history.messages, llm_input]) + http_client = httpx.Client(proxy=f'{proxy.protocol}://{proxy.address}') + try: + self.llm = ChatOpenAI( + model='gpt-4o', + http_client=http_client, + ) + self.llm.tiktoken_model_name = 'gpt-4' + self.llm.temperature = 0.8 + self.llm.top_p = 1 + self.llm.presence_penalty = 0 + chat_history = self.get_chat_history(model_name='gpt-4o') + conversation = RunnableWithMessageHistory( + runnable=self.llm, + get_session_history=lambda _: chat_history, + ) + llm_input = HumanMessage(content=input_content) + if file: + image_count = getattr(self, 'image_count', 0) + input_tokens = self.count_text_tokens([*chat_history.messages, llm_input, *chunks]) + input_tokens += Decimal(image_count) * Decimal('0.13') + else: + input_tokens = self.count_text_tokens([*chat_history.messages, llm_input]) - self.assert_enough_balance(input_tokens, None, model=self.llm.model_name) + self.assert_enough_balance(input_tokens, None, model=self.llm.model_name) - start_time = time.time() - try: + start_time = time.time() if file: input_tokens = self.count_text_tokens([*chat_history.messages]) if sum([len(chunk.content) for chunk in chunks]) > 40_000: @@ -233,6 +234,9 @@ class Raifgpt(Chatgpt): raise ExceededContextLengthError raise + finally: + http_client.close() + output_tokens = self.count_text_tokens([response]) self.logger.info(f'Input количество токенов для raifgpt - {input_tokens}') @@ -257,6 +261,7 @@ class Raifgpt(Chatgpt): def get_pdf_data(self, pdf_data: bytes) -> str: max_batch_size = 3.9 * 1024 * 1024 image_count = 0 + doc = None try: doc = fitz.open(stream=pdf_data, filetype="pdf") raw_texts = {} @@ -268,11 +273,14 @@ class Raifgpt(Chatgpt): if page.get_images(): pages_with_image.append(page_num) if not pages_with_image: - doc.close() + if doc is not None: + doc.close() fitz.TOOLS.store_shrink(100) all_text = "\n".join(raw_texts.get(i, "") for i in sorted(raw_texts)) return all_text if all_text.strip() else "Не удалось извлечь текст из PDF" except Exception as e: + if doc is not None: + doc.close() return f"Ошибка при чтении PDF: {e}" try: batch_images = [] @@ -297,7 +305,8 @@ class Raifgpt(Chatgpt): except Exception: continue if not batch_images: - doc.close() + if doc is not None: + doc.close() fitz.TOOLS.store_shrink(100) return "Не удалось собрать изображения из PDF." batches = [] @@ -345,7 +354,8 @@ class Raifgpt(Chatgpt): if line_text: page_text.append(line_text) ocr_texts[pages[i]] = "\n".join(page_text) - doc.close() + if doc is not None: + doc.close() fitz.TOOLS.store_shrink(100) all_pages = sorted(set(raw_texts) | set(ocr_texts)) final_text = "\n\n".join( @@ -355,6 +365,8 @@ class Raifgpt(Chatgpt): self.image_count = image_count return final_text.strip() or "Не удалось распознать текст" except Exception as e: + if doc is not None: + doc.close() return f"Не удалось обработать файл: {e}" def make_embeddings_prompt(self, document_name: str, section_texts: List[str], question: str) -> str: @@ -66,15 +66,15 @@ class Sora(SimpleService): kind = filetype.guess(input_message.file.read(20)) mime_type = kind.mime if kind else 'application/octet-stream' input_message.file.seek(0) - img = Image.open(input_message.file) - width, height = img.size - required_size = (720, 1280) - current_size = (width, height) - if current_size != required_size: - raise UnsupportedSize(current_size, required_size) - buf = BytesIO() - img.save(buf, format="PNG") - buf.seek(0) + with Image.open(input_message.file) as img: + width, height = img.size + required_size = (720, 1280) + current_size = (width, height) + if current_size != required_size: + raise UnsupportedSize(current_size, required_size) + buf = BytesIO() + img.save(buf, format="PNG") + buf.seek(0) files = { "input_reference": ( input_message.file.name,