llm/wiki_agent.py

from llm.gemini_client import GeminiClient

from llm.utils.wiki_client import WikiClient

SYSTEM_MESSAGE = """you have access to a wikipedia summarizer that can return a summary for a topic. \
Your job is to act as a question answering tool. Whenever you are asked about a question related to knowledge, \
instead of using your internal knowledge (which can be faulty or out of date), \
format a Wikipedia search query string that can help answer the question. \

Remember Wikipedia Entries are usually about a simple entity or event, so keep the \
query short, and about the entity being asked about. Also, don't use your knowledge \
to ask about the answer. Instead form queries about the entity in the question. This \
will help you get the right wikipedia entries for questions when you dont know the answer

### Example 1:

Question: Who won the ICC Cricket World Cup?
Correct Response: Cricket World Cup
Incorrect response: Australia

### Example 2:

Question: Who directed the classic 30s western Stagecoach?
Response: Stagecoach
Incorrect response: John Ford


Below is the question. Return the wikipedia search query you would use \n

### Question:
"""


class WikiSearchAgent:
    def __init__(self):
        self._llm_client = GeminiClient(system_message=SYSTEM_MESSAGE)
        self._wiki_client = WikiClient()

    def get_wikipedia_entry(self, prompt: str) -> str:

        wiki_search_query = self._llm_client.generate_text(prompt)
        wikipedia_page = self._wiki_client.get_pages(wiki_search_query)
        try:
            return wikipedia_page.summary
        except:
            return ""