chatbot.py

import dotenv
from langchain.prompts import PromptTemplate
from langchain_core.runnables import chain
from langchain_chroma import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

from utils import ChromaConstants, save_generated_data

dotenv.load_dotenv()

# - 'Quais docentes possuem expertise em [área específica] com aplicações em [contexto específico]?'
# - 'Pesquisadores que estudam [tópico] e publicaram artigos recentes sobre [detalhe adicional].'
# - 'Quais são as publicações acadêmicas e artigos que estão relacionados a [tópico]?'
# - 'Quais publicações acadêmicas e projetos estão associados ao [docente]?'

user_prompt_refinement_template_str = """
Você é um assistente especializado em análise de perguntas para maximizar a 
recuperação de informações relevantes. Sua tarefa é reescrever a pergunta 
fornecida de forma clara, detalhada e objetiva, destacando os elementos centrais
e estruturando-a para que seja mais compreensível para o sistema de recuperação 
de documentos. Leve em consideração o seguinte:

1. Identifique o problema ou ideia de negócio mencionada na pergunta e sua área 
de aplicação (ex.: agronegócio, IA, infraestrutura).
2. Inclua termos específicos ou técnicos relevantes à área mencionada.
3. Determine se a pergunta envolve colaboração prática, pesquisa teórica ou 
aplicação em uma tecnologia específica.
4. Adicione contexto, se necessário, para que os documentos mais relevantes 
sobre docentes, projetos e artigos sejam priorizados.
5. Identifique, caso seja especificado, o nome de docentes citados.

Reescreva a pergunta final em formato sucinto e priorize a relevância para a 
consulta sobre docentes e seus artigos acadêmicos. Exemplos:

- 'Artigos na área de [área específica] com aplicações em [contexto específico]?'
- 'Pesquisadores que estudam [tópico] e publicaram artigos sobre [detalhe adicional].'
- 'Quais são as publicações acadêmicas que estão relacionados a [tópico]?'
- 'Quais publicações acadêmicas e projetos estão associados ao [docente]?'
- 'Docentes com experiência em projetos de pesquisa aplicada em [tema específico].'
- 'Docentes que colaboram em projetos de [tecnologia/tópico] aplicados a [setor/indústria].'

Caso não seja capaz identificar as informações necessárias para reescrever
a pergunta feita pelo usuário, não faça nenhuma modificação a pergunta original.

Garanta-se de retornar a pergunta otimizada em inglês.

Pergunta original: {question}
Pergunta otimizada:
"""

# Estrutura da resposta:

# - Identificação do problema ou ideia: Resuma a necessidade apresentada pelo agente externo.
# - Conexão com especialistas: Destaque os docentes ou projetos mais relevantes, incluindo
#   suas áreas de atuação e produções relacionadas.
# - Sugestão de aplicação: Explique como o conhecimento ou pesquisa identificado pode ser
#   aplicado para atender à necessidade apresentada.

publication_template_str = """
Você é um assistente especializado em conectar agentes externos, como empresas, 
com os especialistas mais qualificados de uma universidade, com base em suas 
áreas de expertise e produções acadêmicas.
Sua formação combina conhecimento avançado em biblioteconomia, análise de dados 
acadêmicos e comunicação estratégica. Você é capaz de identificar aplicações 
práticas, potenciais de inovação e casos de uso em projetos, dissertações e 
artigos científicos.
Suas diretrizes são as seguintes:

- Baseie-se exclusivamente no contexto fornecido: Use apenas as informações 
  disponíveis no contexto para formular respostas. Não invente dados ou 
  extrapole informações não fornecidas.
- Adapte-se à necessidade do usuário: Interprete claramente o problema ou a 
  ideia apresentada na questão e relacione-a com as áreas de expertise dos 
  docentes mencionados no contexto.
- Seja específico e prático: Proporcione respostas detalhadas, mencionando 
  docentes, projetos ou publicações relevantes e explicando como esses recursos 
  podem atender à necessidade apresentada.
- Se não souber a resposta: Indique educadamente que não há informações suficientes 
  no contexto para responder à pergunta.

Nota importante:

Se a necessidade apresentada for ampla, priorize os especialistas e projetos com
maior relevância. Caso não consiga identificar especialistas adequados no contexto 
fornecido, recomende um direcionamento alternativo.

Garanta-se de retornar a resposta em português independente da língua da questão
ou do contexto. 

Entrada:

- Contexto: 
{context}

- Pergunta: 
{question}
"""


refinement_template = PromptTemplate(
    template=user_prompt_refinement_template_str, input_variables=["question"]
)

prompt_template = PromptTemplate(
    template=publication_template_str, input_variables=["context", "question"]
)

structured_pubs_vectorstore = Chroma(
    persist_directory=ChromaConstants.PUBLICATIONS_CHROMA_PATH,
    embedding_function=OpenAIEmbeddings(),
    collection_name=ChromaConstants.STRUCTURED_COLLECTION,
)

citation_pubs_vectorstore = Chroma(
    persist_directory=ChromaConstants.PUBLICATIONS_CHROMA_PATH,
    embedding_function=OpenAIEmbeddings(),
    collection_name=ChromaConstants.CITATION_COLLECTION,
)


def get_structured_chain(k):
    """Create a new function to answer questions using `model` and `k` documents."""
    chat_model = ChatOpenAI(model="gpt-4o-2024-08-06", temperature=0)

    @chain
    def pubs_retriever(query):
        """Create a retriver that include the score for each document retrieved in the answer."""
        docs, scores = zip(
            *structured_pubs_vectorstore.similarity_search_with_score(query, k=k)
        )
        for doc, score in zip(docs, scores):
            doc.metadata["score"] = score

        return docs

    def get_assistant_answer(question, question_id=None, question_dir=None):
        """Answer for the question provided using research articles as context."""
        question_refinement_chain = refinement_template | chat_model | StrOutputParser()
        publications_chain = prompt_template | chat_model | StrOutputParser()

        context = pubs_retriever.invoke(question)
        context_content = [doc.page_content for doc in context]
        context_str = "\n\n".join(context_content)

        refined_question = question_refinement_chain.invoke(question)
        refined_context = pubs_retriever.invoke(refined_question)
        refined_context_content = [doc.page_content for doc in refined_context]
        refined_context_str = "\n\n".join(refined_context_content)

        answer = publications_chain.invoke(
            {"context": context_str, "question": question}
        )
        refined_answer = publications_chain.invoke(
            {"context": refined_context_str, "question": question}
        )

        if question_id is not None and question_dir is not None:
            save_generated_data(
                "structured-data",
                question,
                question_id,
                question_dir,
                refined_question,
                context,
                refined_context,
                answer,
                refined_answer,
                {"k": k},
            )

        return refined_answer, refined_question

    return get_assistant_answer


def get_citation_chain(k):
    """Create a new function to answer questions using `model` and `k` documents."""
    chat_model = ChatOpenAI(model="gpt-4o-2024-08-06", temperature=0)

    @chain
    def pubs_retriever(query):
        """Create a retriver that include the score for each document retrieved in the answer."""
        docs, scores = zip(
            *citation_pubs_vectorstore.similarity_search_with_score(query, k=k)
        )
        for doc, score in zip(docs, scores):
            doc.metadata["score"] = score

        return docs

    def get_assistant_answer(question, question_id=None, question_dir=None):
        """Answer for the question provided using research articles as context."""
        question_refinement_chain = refinement_template | chat_model | StrOutputParser()
        publications_chain = prompt_template | chat_model | StrOutputParser()

        context = pubs_retriever.invoke(question)
        context_content = [doc.page_content for doc in context]
        context_str = "\n\n".join(context_content)

        refined_question = question_refinement_chain.invoke(question)
        refined_context = pubs_retriever.invoke(refined_question)
        refined_context_content = [doc.page_content for doc in refined_context]
        refined_context_str = "\n\n".join(refined_context_content)

        answer = publications_chain.invoke(
            {"context": context_str, "question": question}
        )
        refined_answer = publications_chain.invoke(
            {"context": refined_context_str, "question": question}
        )

        if question_id is not None and question_dir is not None:
            save_generated_data(
                "citation-data",
                question,
                question_id,
                question_dir,
                refined_question,
                context,
                refined_context,
                answer,
                refined_answer,
                {"k": k},
            )

        return refined_answer, refined_question

    return get_assistant_answer