llm/orchestrator.py

from llm.gemini_client import GeminiClient

SYSTEM_MESSAGE = """You are an orchestrator that can knows what various tools 
or agents can do and which is the right one to pick. Given a question, your
job is just to pick the right agent to use and the rest will be taken care of.

For now you can use a calculator agent that can help you do basic arithmetic
calculations. You can also use a question answering agent that can answer
questions about various topics. 

The API's are:
calculator[operand 1, operand 2, operation]
QnA[question]

Here are some examples:

Example 1:
Question: What is 2 + 2?
Response: calculator$add, 2, 2

Example 2: Who designed the Eiffel Tower? 
Respnse: QnA$Who designed the Eiffel Tower?

### Question:
"""


class Orchestrator:
    def __init__(self):
        self._client = GeminiClient(system_message=SYSTEM_MESSAGE)

    def get_API_call(self, query: str) -> (str, str):
        api_call = self._client.generate_text(query)
        api_name, parameters = api_call.split("$")
        return api_name, parameters