llm/qa_agent.py

from llm.gemini_client import GeminiClient
from llm.wiki_agent import WikiSearchAgent

SYSTEM_MESSAGE = """You are a Question Answering tool that can answer various 
trivia questions. However, you might be asked questions that is beyond your
knowledge or recent events that you might not be trained on 
(beyond training cutoff). So, if there is Wikipedia page entry provided,
use that to answer the question. Just return the answer, don't make a 
verbose response."""


class QnAAgent:
    def __init__(self):
        self._client = GeminiClient(system_message=SYSTEM_MESSAGE)
        self._wiki_tool = WikiSearchAgent()

    @staticmethod
    def _format_prompt(query: str, wiki_page: str) -> str:
        return f"\n###Question:{query} \
        \n###Wikipedia Page:{wiki_page}"

    def get_answer(self, query: str, use_context: bool = True) -> [str, str]:
        if use_context:
            wiki_page = self._wiki_tool.get_wikipedia_entry(query)
            prompt = self._format_prompt(query, wiki_page)
        else:
            wiki_page = ""
            prompt = query
        return self._client.generate_text(prompt), wiki_page