src/research_assistant/components/solver.py

from typing import List, Optional

from langchain_anthropic.chat_models import ChatAnthropic
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field

from research_assistant.app_logging import app_logger
from research_assistant.constants import HEILMEIER_CATECHISM


class SolverResponse(BaseModel):
    """The input to the summarizer tool function."""

    answer: str = Field(
        ...,
        description="The summary of the research article ",
    )


def get_solver(llm: ChatAnthropic):
    prompt_template = ChatPromptTemplate.from_messages(
        [
            SystemMessage(
                (
                    "You are a research assistant responsible for simplifying and explaining the core concepts of a research article to a user who may not be familiar with technical terms. You will be given a plan created by a planner tool, which breaks down the main ideas of the research article through a series of questions. Each question is paired with an answer, providing insights into the purpose, methodology, and key findings of the article. Your task is to synthesize these questions and answers to produce a clear, concise summary that captures the main message of the research article. The summary should enable the user to understand the article's significance, its contributions, and whether it contains information relevant to their needs or goals. Your summary should be:\n"
                    "1. Simple and accessible, avoiding technical jargon.\n"
                    "2. Comprehensive enough to convey the article's goals and key insights.\n"
                    "3. Informative so that the user can decide if reading the full article is worth their time.\n"
                    "Please proceed by summarizing based on these questions and answers. Make sure to respond in the markdown format."
                ),
            ),
            MessagesPlaceholder(variable_name="context", optional=True),
            MessagesPlaceholder(variable_name="text"),
        ]
    )
    solver = prompt_template | llm

    def get_joined_answer(
        input: str,
        _context: Optional[List[str]] = None,  # TODO: rename when context is used
        _config: Optional[RunnableConfig] = None,  # TODO: rename when config is used
    ):
        response = solver.invoke(
            {
                "text": [
                    HumanMessage(
                        (
                            f"Solve the following task or question. To solve the question, we have made a step-by-step plan and retrieved corresponding evidence for each plan. Use them with caution since long evidence might contain irrelevant information. Here's the plan with the evidence:\n{input}\n\n"
                            f"Now solve the question or task according to the provided evidence above. Respond with:\n{HEILMEIER_CATECHISM}\n\n"
                            "Since you are a research assistant, you need to be as detailed as possible to help me understand your breakdown of the research document. Assume I have no prior knowledge of the document’s content. Make it clear, comprehensive, and easy to understand. Avoid complex language and technical jargon. The more in-depth the explanation, the better\n"
                            "The output should be answering all the Heilmeier catechism questions using the obtained evidence information. Answer all the questions.\n"
                        )
                    )
                ]
            }
        )
        if not isinstance(response.content, str):
            app_logger.info(
                "The response from solver is not a string. It is %s",
                type(response.content),
            )
            response.content = str(response.content)
        return SolverResponse(answer=response.content)

    return StructuredTool.from_function(
        name="solver",
        func=get_joined_answer,
        description="This is a tool which takes in a list of questions and the answers provide to those questions, and generates a summary by using all of this information.",
    )