Spaces:
Sleeping
Sleeping
File size: 3,628 Bytes
e3e865e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
from llama_index.core.workflow import Context
import os
from typing import List
from args import LLMInterface
from llm_factory import LLMFactory
from llama_index.core.agent.workflow import AgentWorkflow
class IAgent():
def __init__(self, temperature, max_tokens, sys_prompt_file, llm_itf: LLMInterface):
self.temperature, self.max_tokens = temperature, max_tokens
# Load the system prompt from a file
system_prompt_path = os.path.join(os.getcwd(), "system_prompts", sys_prompt_file)
self.system_prompt = ""
with open(system_prompt_path, "r") as file:
self.system_prompt = file.read().strip()
# Initialize the tool agents
self.tools = self.setup_tools()
self.slaves: List[IAgent] = self.setup_slaves()
# Define the LLM and agent
self.llm = LLMFactory.create(llm_itf, self.system_prompt, temperature, max_tokens)
self.agent = self._setup_agent()
self.ctx = Context(self.agent)
def setup_tools(self) -> List:
"""
Abstract method to set up the tools.
This method must be overridden by subclasses to define custom tools this agent can use.
"""
raise NotImplementedError("Subclasses must implement the setup_tools method.")
def setup_slaves(self) -> List:
"""
Abstract method to set up the slave agents.
This method must be overridden by subclasses to define custom sub-agents this agent can use.
"""
raise NotImplementedError("Subclasses must implement the setup_slaves method.")
def _setup_agent(self) -> AgentWorkflow:
"""
Initializes and returns an agent workflow based on the presence of tools and slaves.
If both `self.tools` and `self.slaves` are empty, it sets up a default agent using the provided language model (`self.llm`).
Otherwise, it creates an agent workflow using the combined list of tools and slaves with the language model.
Returns:
AgentWorkflow: An instance of the agent workflow configured with the appropriate tools and language model.
"""
if not self.tools and not self.slaves:
return AgentWorkflow.setup_agent(llm=self.llm)
# Create tools from slaves: each tool calls slave.query(question) asynchronously
slave_tools = []
for slave in self.slaves:
slave_tools.append(slave.query)
self.tools.extend(slave_tools)
return AgentWorkflow.from_tools_or_functions(
self.tools,
llm=self.llm
)
def get_system_prompt(self) -> str:
"""
Retrieves the system prompt.
Returns:
str: The system prompt string.
"""
return self.system_prompt
async def query(self, question: str) -> str:
"""
Asynchronously queries the agent with a given question and returns the response.
Args:
question (str): The question to be sent to the agent.
Returns:
str: The response from the agent as a string.
"""
response = await self.agent.run(question, ctx=self.ctx)
response = str(response)
return response
def clear_context(self):
"""
Clears the current context of the agent, resetting any conversation history.
This is useful when starting a new conversation or when the context needs to be refreshed.
"""
self.ctx = Context(self.agent)
if not self.slaves:
return
for slave in self.slaves:
slave.clear_context()
|