rag.py

from db import db
import logging
from summary import file_summary

logger = logging.getLogger("rag")
logging.basicConfig(
    format="%(asctime)s %(levelname)-8s %(message)s",
    level=logging.INFO,
    datefmt="%Y-%m-%d %H:%M:%S",
)


message_template = """\
Далее представлена информацию по опыту нашей компании
---------------------
{retrieved_chunks}
---------------------
Далее представлен запрос потенциального проекта
---------------------
{request_content}
---------------------
При проведении анализа опирайся только на представленную информацию"""

# Функция для обработки запроса к LLM
def process_query(req_file, system_prompt, llm, temperature, alpha):
    logger.info("Process query")

    req_file_content = file_summary(req_file)

    logger.info("Retrive docs")

    docs = db.query(req_file_content, top_k=3, alpha=alpha)
    
    logger.info(f"Retrived {len(docs['ids'])} docs")

    doc_context = '\n\n'.join(docs['documents'])

    # Создание контекста из файлов
    user_message = message_template.format(retrieved_chunks=doc_context, request_content=req_file_content)

    # Формирование сообщения для LLM
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_message}
    ]

    response = llm.chat(messages, temperature)

    # Получение ответа от LLM
    llm_response = response.choices[0].message.content
    return llm_response