
from services.chat_service import ChatService
from Copilot.LLMLayer.Prompts.document_prompts import *
from DocumentManager.DocumentsManager import DocumentManager
from Copilot.LLMLayer.Utils.LLMUtils import llm_json_dumps
import json

class RetrievalService:
    def __init__(self, doc_manager:DocumentManager, text_model_name):
        self.doc_manager = doc_manager
        self.text_model_name = text_model_name

    def build_query(self, incident:str, user_message):
        incident.pop("parameter_histories", None)
        incident_context = llm_json_dumps(incident)
        query_prompt = f"""
        Incident:
        {incident_context}

        User observation:
        {user_message}

        Generate search query:
        """

        query, ret = ChatService.ask_openai(
            query_prompt,
            history=[],
            text_model_name=self.text_model_name,
            system_prompt=QUERY_BUILDER_PROMPT
        )

        return query if ret else None
    
    def has_any_dcoument(self,):
        return self.doc_manager.has_index()

    def retrieve(self, query):
        if not query or not self.doc_manager.has_index():
            return ""

        chunks, _ = self.doc_manager.search(query, top_k=10)
        return ChatService.chunks_to_prompt(chunks)