Swarm_Ex - Instructor Integration

Mix.install(
  [
    {:instructor, "~> 0.0.5"},
    {:openai_ex, "~> 0.8.4"},
    {:swarm_ex, git: "https://github.com/nrrso/swarm_ex.git", branch: "main"},
    {:jason, "~> 1.4"}
  ],
  config: [
    instructor: [
      adapter: Instructor.Adapters.OpenAI,
      openai: [api_key: System.fetch_env!("LB_OPENAI_API_KEY")]
    ]
  ]
)

Define Example Modules

defmodule StarWarsClassifier do
  use Ecto.Schema
  use Instructor.Validator

  @doc """
  ## Field Descriptions:
  - is_star_wars: Whether or not the message is related to star wars
  - reason: A short, less than 10 word rationalization for the classification.
  - score: A confidence score between 0.0 and 1.0 for the classification.
  """
  @primary_key false
  embedded_schema do
    field(:is_star_wars, :boolean)
    field(:reason, :string)
    field(:score, :float)
  end

  @impl true
  def validate_changeset(changeset) do
    changeset
    |> Ecto.Changeset.validate_number(:score,
      greater_than_or_equal_to: 0.0,
      less_than_or_equal_to: 1.0
    )
  end

  def is_star_wars?(text) do
    Instructor.chat_completion(
      model: "gpt-4o-mini",
      response_model: __MODULE__,
      max_retries: 3,
      messages: [
        %{
          role: "user",
          content: """
          Your purpose is to classify if an incoming message is related to star wars or not.
  
          Classify the following message: 
          ```
          #{text}
          ```
          """
        }
      ]
    )
  end
end

StarWarsClassifier.is_star_wars?("Christmas cookies are the best!")

StarWarsClassifier.is_star_wars?("I have a bad feeling about this.")

defmodule YodaReply do
  use Ecto.Schema
  use Instructor.Validator

  @doc """
  ## Field Descriptions:
  - quote: A fitting reply in the tone and style of Yoda.
  """
  @primary_key false
  embedded_schema do
    field(:reply, :string)
  end

  def yoda_reply?(text) do
    Instructor.chat_completion(
      model: "gpt-4o-mini",
      response_model: __MODULE__,
      max_retries: 3,
      messages: [
        %{
          role: "user",
          content: """
          Your purpose is to reply in the tone and style of Master Yoda.
  
          Reply to the following message: 
          ```
          #{text}
          ```
          """
        }
      ]
    )
  end
end

# Message Classification Tool
defmodule ClassifyMessageTool do
  @behaviour SwarmEx.Tool

  @impl true
  def execute(%{message: message}) do
    # For now, simple keyword matching
    {:ok, %{is_star_wars: _is_star_wars}} = StarWarsClassifier.is_star_wars?(message)
  end

  @impl true
  def validate(%{message: _}), do: :ok
  def validate(_), do: {:error, :invalid_args}

  @impl true
  def cleanup(_), do: :ok
end

# Star Wars Response Tool
defmodule StarWarsResponseTool do
  @behaviour SwarmEx.Tool

  @impl true
  def execute(%{message: message}) do
    {:ok, %{reply: _reply}} = YodaReply.yoda_reply?(message)
  end

  @impl true
  def validate(%{message: _}), do: :ok
  def validate(_), do: {:error, :invalid_args}

  @impl true
  def cleanup(_), do: :ok
end

defmodule TriageAgent do
  use SwarmEx.Agent

  @impl true
  def init(opts) do
    {:ok, opts}
  end

  @impl true
  def terminate(_reason, _state), do: :ok

  @impl true
  def handle_message(message, state) when is_binary(message) do
    # Directly execute the ClassifyMessageTool instead of using execute_tool
    case ClassifyMessageTool.execute(%{message: message}) do
      {:ok, %{is_star_wars: true}} ->
        Logger.info("Message classified as Star Wars related. Handing off to StarWarsAgent.")
        handoff_to_star_wars(message, state)
      
      {:ok, %{is_star_wars: false}} ->
        Logger.info("Message not Star Wars related. Providing general response.")
        {:ok, "Hi, how can I help?", state}
      
      {:error, reason} ->
        Logger.error("Classification failed: #{inspect(reason)}")
        {:error, "Sorry, I couldn't process your message."}
    end
  end

  @impl true
  def handle_tool(:respond, _args, state) do
    {:ok, state}
  end

  defp handoff_to_star_wars(message, state) do
    case state[:star_wars_agent] do
      nil -> 
        {:error, "Star Wars agent not configured"}
      agent_pid -> 
        case SwarmEx.send_message_to_pid(agent_pid, message) do
          {:ok, response} -> {:ok, response, state}
          error -> error
        end
    end
  end
end

# Star Wars Specialist Agent
defmodule StarWarsAgent do
  use SwarmEx.Agent

  @impl true
  def init(opts) do
    {:ok, opts}
  end

  @impl true
  def terminate(_reason, _state), do: :ok

  @impl true
  def handle_message(message, state) when is_binary(message) do
    case StarWarsResponseTool.execute(%{message: message}) do
      {:ok, %{reply: reply}} -> 
        Logger.info("Generated a response from Yoda.")
        {:ok, reply, state}
      {:error, reason} -> 
        Logger.error("Reply failed: #{inspect(reason)}")
        {:error, "Sorry, I couldn't reply to your message."}
    end
  end

  @impl true
  def handle_tool(:respond, _args, state) do
    {:ok, state}
  end
end

Instantiate Tools and Agents

# Create network and agents
{:ok, network} = SwarmEx.create_network()

# Register our tools
SwarmEx.register_tool(ClassifyMessageTool)
SwarmEx.register_tool(StarWarsResponseTool)

# Create the Star Wars agent first
{:ok, star_wars_agent} = SwarmEx.create_agent(network, StarWarsAgent, name: "star_wars_agent")

SwarmEx.Client.list_agents(network)

# Create the triage agent with reference to the Star Wars agent
{:ok, triage_agent} = SwarmEx.create_agent(
  network, 
  TriageAgent, 
  name: "triage_agent",
  star_wars_agent: star_wars_agent
)

SwarmEx.Client.list_agents(network)

Sending Message

# Test with non-Star Wars message
{:ok, regular_response} = SwarmEx.send_message(network, "triage_agent", "Hello, what's the weather like?")
IO.puts("Regular response: #{regular_response}")

# Test with Star Wars message
{:ok, star_wars_response} = SwarmEx.send_message(network, "triage_agent", "Tell me about the Jedi")
IO.puts("Star Wars response: #{star_wars_response}")