/ src / python / txtai / agent / base.py
base.py
  1  """
  2  Agent module
  3  """
  4  
  5  import os
  6  
  7  from collections import deque
  8  
  9  from jinja2 import Template
 10  
 11  from .factory import ProcessFactory
 12  
 13  
 14  class Agent:
 15      """
 16      An agent automatically creates workflows to answer multi-faceted user requests. Agents iteratively prompt and/or interface with tools to
 17      step through a process and ultimately come to an answer for a request.
 18  
 19      Agents excel at complex tasks where multiple tools and/or methods are required. They incorporate a level of randomness similar to different
 20      people working on the same task. When the request is simple and/or there is a rule-based process, other methods such as RAG and Workflows
 21      should be explored.
 22      """
 23  
 24      def __init__(self, template=None, memory=None, **kwargs):
 25          """
 26          Creates a new Agent.
 27  
 28          Args:
 29              template: optional prompt jinja template, must include {{ text }} and {{ memory }} placeholders
 30              memory: number of prior outputs to keep as "memory", defaults to None for no memory
 31              kwargs: arguments to pass to the underlying Agent backend and LLM pipeline instance
 32          """
 33  
 34          # Ensure backwards compatibility
 35          if "max_iterations" in kwargs:
 36              kwargs["max_steps"] = kwargs.pop("max_iterations")
 37  
 38          # Custom instructions
 39          if "instructions" in kwargs:
 40              kwargs["instructions"] = self.instructions(kwargs)
 41  
 42          # Create agent process runner
 43          self.process = ProcessFactory.create(kwargs)
 44  
 45          # Tools dictionary
 46          self.tools = self.process.tools
 47  
 48          # Agent memory
 49          self.memory = {}
 50          self.window = memory
 51  
 52          # Create template
 53          self.template = Template(
 54              template
 55              if template
 56              else """{{ text }}
 57  {% if memory %}
 58  Use the following conversation history to help answer the question above.
 59  
 60  {{ memory }}
 61  
 62  If the history is irrelevant, forget it and use other tools to answer the question.
 63  {% endif %}
 64  """
 65          )
 66  
 67      def __call__(self, text, maxlength=8192, stream=False, session=None, reset=False, **kwargs):
 68          """
 69          Runs an agent loop.
 70  
 71          Args:
 72              text: instructions to run
 73              maxlength: maximum sequence length
 74              stream: stream response if True, defaults to False
 75              session: session id for stored memory, defaults to None which shares all memory
 76              reset: clears previously stored memory if True, defaults to False
 77              kwargs: additional keyword arguments
 78  
 79          Returns:
 80              result
 81          """
 82  
 83          # Process parameters
 84          self.process.model.parameters(maxlength)
 85  
 86          # Create memory, if necessary
 87          if self.window and (session not in self.memory or reset):
 88              self.memory[session] = deque(maxlen=self.window)
 89  
 90          # Run agent loop
 91          output = self.process.run(self.prompt(text, session), stream=stream, **kwargs)
 92  
 93          # Add output to memory, if necessary
 94          if session in self.memory:
 95              self.memory[session].append((text, output))
 96  
 97          return output
 98  
 99      def prompt(self, text, session):
100          """
101          Generates the full agent prompt using the input instructions. Adds in agent memory
102          if available.
103  
104          Args:
105              text: instructions to run
106              session: session id for stored memory
107  
108          Returns:
109              formatted instructions
110          """
111  
112          # pylint: disable=E1133
113          memory = []
114          if self.memory.get(session):
115              # Add memory context
116              for request, output in self.memory[session]:
117                  memory.append(f"User: {request}\nAssistant: {output}")
118  
119              memory = "\n\n".join(memory)
120  
121          # Command template with memory
122          return self.template.render(text=text, memory=memory)
123  
124      def instructions(self, config):
125          """
126          Reads and formats custom instructions. Supports agents.md files.
127  
128          Args:
129              config: agent configuration
130  
131          Returns:
132              agent instructions, if any
133          """
134  
135          # Check if this is a file path (i.e. agents.md)
136          path = config.pop("instructions", None)
137          if path and os.path.isfile(path):
138              with open(path, encoding="utf-8") as f:
139                  # Read entire file
140                  return f"\nBelow is a set of general instructions to follow:\n\n{f.read()}"
141  
142          return path