|
|
from magentic import FunctionCall, OpenaiChatModel, prompt_chain |
|
|
import requests |
|
|
|
|
|
from BasicAgent import BasicAgent |
|
|
|
|
|
import os |
|
|
|
|
|
|
|
|
class CustomAgent(BasicAgent): |
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
self._system_prompt = """\ |
|
|
You are a general AI assistant. I will ask you a question. Provide an answer like so: YOUR FINAL ANSWER. DO NOT SAY ANYTHING EXCEPT YOUR FINAL ANSWER. |
|
|
Only use the provided tools if your own knowledge would be insufficient. If you are unable to answer the question using the provided tools, give your best guess. |
|
|
Your final answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. |
|
|
If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. |
|
|
If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. |
|
|
If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. |
|
|
Question: {question} |
|
|
""" |
|
|
self._model = OpenaiChatModel( |
|
|
model="llama-3.1-8b-instant", |
|
|
api_key=os.environ["GROQ_API_KEY"], |
|
|
base_url="https://api.groq.com/openai/v1/", |
|
|
temperature=0.7, |
|
|
max_tokens=512 |
|
|
) |
|
|
print("CustomAgent initialized.") |
|
|
|
|
|
def __call__(self, question: str) -> str: |
|
|
""" |
|
|
this is where the magic happens |
|
|
""" |
|
|
print(f"Agent received question (first 50 chars): {question[:50]}...") |
|
|
@prompt_chain( |
|
|
template=self._system_prompt, |
|
|
model=self._model, |
|
|
functions=[self.websearch], |
|
|
) |
|
|
def _magic(question: str) -> FunctionCall | str: |
|
|
""" |
|
|
An agent, using websearch to gather information and |
|
|
answer a user's query. |
|
|
""" |
|
|
... |
|
|
result = _magic(question) |
|
|
print("Agent returning actual answer:", result) |
|
|
return result |
|
|
|
|
|
@staticmethod |
|
|
def websearch(query: str) -> str: |
|
|
""" |
|
|
Performs a web search using an external web search API. Note that better results may be obtained from varying the wording of the query. |
|
|
Args: |
|
|
query (str): The search query string. |
|
|
Returns: |
|
|
str: The response from the web search API as a string. |
|
|
""" |
|
|
try: |
|
|
r = requests.post( |
|
|
f"{os.environ['WEBSEARCH_API_URL']}/tools/websearch", |
|
|
json={"query": query} |
|
|
) |
|
|
r.raise_for_status() |
|
|
response = r.text |
|
|
except Exception as e: |
|
|
print("Error:", e) |
|
|
response = "An error has occurred while using websearch, please inform the user of this" |
|
|
print(f"Agent searched for {query}. Response: {response}") |
|
|
|
