dspy-agent-state-machine-dry.py

import dspy
from dspy.functional import TypedPredictor
import os
from dotenv import load_dotenv
from transitions import Machine

# Load the OpenAI API key from the .env file
load_dotenv()

# Create a language model using the OpenAI API
llm = dspy.OpenAI(
    model='gpt-4',
    api_key=os.environ['OPENAI_API_KEY'],
    max_tokens=400
)

# Configure the dspy settings to use the language model
dspy.settings.configure(lm=llm)

class DecisionSignature(dspy.Signature):
    input_text = dspy.InputField(desc="The input text to be processed")
    rationale = dspy.OutputField(desc="The rationale for the decision")
    decision: bool = dspy.OutputField(desc="True if the input text contains the final answer, False otherwise")

class Agent(Machine):
    states = ['start', 'thought', 'acted', 'observed', 'concluded']

    def __init__(self, llm, objective=None):
        super().__init__(states=Agent.states, initial='start')
        self.llm = llm
        self.objective = objective
        self.memory = []
        self.add_transition(trigger='think', source='start', dest='thought')
        self.add_transition(trigger='act', source='thought', dest='acted')
        self.add_transition(trigger='observe', source='acted', dest='observed')
        self.add_transition(trigger='decide', source='observed', dest='concluded')

    def process_step(self, action, message, next_state):
        print(action + "...")
        prompt = f"{message}: {self.objective} {' '.join(self.memory)}"
        response = self.llm(prompt).pop()
        self.memory.append(response)
        print(response)
        self.state = next_state

    def decide(self):
        print("Deciding...")
        prompt = "Based on your observations, make a decision: " + ' '.join(self.memory)
        decision_maker = TypedPredictor(DecisionSignature)
        response = decision_maker(input_text=prompt)

        if response.decision:
            self.state = 'concluded'
            final_answer = self.llm(f"What is the final answer to this: {self.objective}, given this: {' '.join(self.memory)}").pop()
            print("The final answer is: " + final_answer)
            return final_answer
        else:
            self.memory.append("Decision not reached because " + response.rationale)
            self.state = 'start'

    def execute(self):
        actions = {
            'start': (self.process_step, 'Thinking', "Think step by step about how to correctly answer this", 'thought'),
            'thought': (self.process_step, 'Acting', "Execute the thinking based on the information you have", 'acted'),
            'acted': (self.process_step, 'Observing', "Analyze the results of your actions", 'observed'),
            'observed': (self.decide,)
        }
        while self.state != 'concluded':
            action = actions[self.state]
            if len(action) == 4:
                action[0](action[1], action[2], action[3])
            else:
                action[0]()

agent = Agent(llm, objective="What is the double of the sum of Barack Obama and his wife's age in April 2024?")
agent.execute()