How to define a custom output including data from OpenAI functions?

Question

i am trying to create a chat assistent using Langchain and OpenAI functions, that has multiple agents to generate data to reply user questions.

In the example below, the agent is calling a api to find location data and generates a response like { output: "some answer string with markups" }.

Now i want to change the output to a format like: { output: "some answer string", docs: [array of documents returned from the function tool]

Any idea how to change the code below to get that done? Thanks for your help!

server/router.ts

import express, { Request, Response } from 'express'
import { ChatOpenAI } from 'langchain/chat_models/openai'
import { AgentExecutor } from 'langchain/agents'
import { ChatPromptTemplate, MessagesPlaceholder } from 'langchain/prompts'
import { formatToOpenAIFunction } from 'langchain/tools'
import { RunnableSequence } from 'langchain/schema/runnable'
import { OpenAIFunctionsAgentOutputParser } from 'langchain/agents/openai/output_parser'
import { StructuredOutputParser } from 'langchain/output_parsers'
import {
  AIMessage,
  AgentStep,
  BaseMessage,
  FunctionMessage,
} from 'langchain/schema'
import tools from './tools'

const router: express.Router = express.Router()

router.post('/', async (req: Request, res: Response) => {
  try {
    const { messages } = req.body
    const input = messages[messages.length - 1].content

    const model = new ChatOpenAI({
      openAIApiKey: 'XYZ',
      modelName: 'gpt-4',
      temperature: 0,
    })

    const prompt = ChatPromptTemplate.fromMessages([
      [
        'ai',
        `Du bist ein hilfsbereiter Assistent bei der Planung einer Party.
        Wenn du das Location Tool nutzt und Locations gefunden hast, fasse den Namen, Adresse und Telefonnummer zusammen.
        `,
      ],
      ['human', '{input}'],
      new MessagesPlaceholder('agent_scratchpad'),
    ])

    const modelWithFunctions = model.bind({
      functions: [...tools.map((tool) => formatToOpenAIFunction(tool))],
    })

    const formatAgentSteps = (steps: AgentStep[]): BaseMessage[] =>
      steps.flatMap(({ action, observation }) => {
        if ('messageLog' in action && action.messageLog !== undefined) {
          const log = action.messageLog as BaseMessage[]
          return log.concat(new FunctionMessage(observation, action.tool))
        }
        return [new AIMessage(action.log)]
      })

    const runnableAgent = RunnableSequence.from([
      {
        input: (i: { input: string; steps: AgentStep[] }) => i.input,
        agent_scratchpad: (i: { input: string; steps: AgentStep[] }) =>
          formatAgentSteps(i.steps),
      },
      prompt,
      modelWithFunctions,
      new OpenAIFunctionsAgentOutputParser(),
    ])

    const executor = AgentExecutor.fromAgentAndTools({
      agent: runnableAgent,
      tools,
    })

    const response = await executor.invoke({
      input,
    })  // <=  returns: { output: "some answer string with markups" }

    res.json(response) 
  } catch (e) {
    console.log(e)
    res.status(500).json({ error: e.message })
  }
})

server/tools/index.ts

import LocationSearchTool from './LocationSearchTool'

export default [new LocationSearchTool()]

tools/Locationsearchtool.ts

import payload from 'payload'
import Ajv from 'ajv'
import { z } from 'zod'
import { ToolParams, StructuredTool } from 'langchain/tools'

type LocationSearchToolInput = {
  city: string
  range: number
  lat: string
  lng: string
}

class LocationSearchTool extends StructuredTool {

  static lc_name() {
    return 'LocationSearchTool'
  }

  name = 'location-search-tool'

  description = `Dieses Tool performed eine Suche nach der passenden Location. 
                        Die Aktionseingabe sollte diesem JSON-Schema entsprechen:
                        {{"city":{{"type":"string","description":"die Stadt in dessen Umkreis die Location liegen sollte"}},
                        "range":{{"type":"integer","description":"der Umkreis um die Stadt in der gesucht werden soll"}}}}
                        Wenn keine Stadt angegeben wird, frage einfach nach der Stadt und dem Umkreis.
                        Das Aktionsergebnis entspricht diesem JSON-Schema: 
                        {{"locations":{{"type":"array","items":{{"type":"string"}},"description":"die Liste der gefundenen Locations"}}}}
                        Du beantwortest die Frage auf Deutsch. Der Text sollte sich auf die Stadt und den Umkreis beziehen und die gefundenen Locations auflisten. Wenn es die Stadt mehrfach mit dem gleichen Namen gibt, frage nach der Postleitzahl.
                        `

  ajv = new Ajv({ strict: false })

  inputSchema = {
    city: {
      type: 'string',
      description: 'die Stadt in dessen Umkreis die Location liegen sollte',
    },
    range: {
      type: 'integer',
      description:
        'der Umkreis in kilometer um die Stadt in der gesucht werden soll',
    },
  }

  outputSchema = {
    locations: {
      type: 'array',
      items: {
        type: 'string',
      },
      description: 'Liste mit den Namen der gefundenen Locations',
    },
    docs: {
      type: 'array',
      description: 'Liste mit den Namen der gefundenen Locations',
      items: {
        type: 'object',
        description: 'Objekt mit den Daten der gefundenen Location',
        properties: {
          name: {
            type: 'string',
            description: 'Name der Location',
          },
          street: {
            type: 'string',
            description: 'Straße der Location',
          },
          streetNumber: {
            type: 'string',
            description: 'Hausnummer der Location',
          },
          postalCode: {
            type: 'string',
            description: 'Postleitzahl der Location',
          },
          city: {
            type: 'string',
            description: 'Stadt der Location',
          },
          country: {
            type: 'string',
            description: 'Land der Location',
          },
          phone: {
            type: 'string',
            description: 'Telefonnummer der Location',
          },
          website: {
            type: 'string',
            description: 'Website der Location',
          },
          slug: {
            type: 'string',
            description: 'Slug der Location',
          },
        },
      },
    },
  }

  schema = z.object({
    city: z
      .string()
      .describe('die Stadt in dessen Umkreis die Location liegen sollte'),
    lat: z
      .string()
      .describe(
        'der latitude der Stadt in dessen Umkreis die Location liegen sollte',
      ),
    lng: z
      .string()
      .describe(
        'der longitude der Stadt in dessen Umkreis die Location liegen sollte',
      ),
    range: z
      .number()
      .describe(
        'der Umkreis in kilometer um die Stadt in der gesucht werden soll',
      ),
  })

  validate(data, schema) {
    if (schema) {
      const validateSchema = this.ajv.compile(schema)
      if (!validateSchema(data)) {
        throw new Error(this.ajv.errorsText(validateSchema.errors))
      }
    }
  }

  constructor(params?: ToolParams) {
    super(params)
  }

  async _call(input: LocationSearchToolInput) {

    let output
    try {
      // const input = JSON.parse(arg)
      this.validate(input, this.inputSchema)

      const { city, range, lat, lng } = input
      const rangeInMeters = range * 1000
      const results = await payload.find({
        collection: 'business-profiles',
        where: {
          category: { in: ['6489ee1388b6dc4236dfdcde'] },
          location: { near: [lng, lat, rangeInMeters, 0] },
        },
      })

      const locations = results.docs.map((result) => result.name)

      output = {
        locations,
        docs: results.docs,
      }
      this.validate(output, this.outputSchema)
    } catch (err) {
      output = { error: err.message || err }
    }
 
    return JSON.stringify(output)
  }
}

export default LocationSearchTool