import { OpenAIAgent } from "llamaindex";
import { tools } from "./tools";

const agent = new OpenAIAgent({
  tools: [...tools],
});
const { response } = await agent.chat({
  message: "What is weather today?",
  stream: false,
});

console.log(response.message.content);

import { AgentRunner, OpenAI } from "llamaindex";

class MyLLM extends OpenAI {}

export class MyAgentWorker extends AgentWorker<MyLLM> {
  taskHandler = MyAgent.taskHandler;
}

export class MyAgent extends AgentRunner<MyLLM> {
  constructor(params: Params) {
    super({
      llm: params.llm,
      chatHistory: params.chatHistory ?? [],
      systemPrompt: params.systemPrompt ?? null,
      runner: new MyAgentWorker(),
      tools:
        "tools" in params
          ? params.tools
          : params.toolRetriever.retrieve.bind(params.toolRetriever),
    });
  }

  // create store is a function to create a store for each task, by default it only includes `messages` and `toolOutputs`
  createStore = AgentRunner.defaultCreateStore;

  static taskHandler: TaskHandler<Anthropic> = async (step, enqueueOutput) => {
    const { llm, stream } = step.context;
    // initialize the input
    const response = await llm.chat({
      stream,
      messages: step.context.store.messages,
    });
    // store the response for next task step
    step.context.store.messages = [
      ...step.context.store.messages,
      response.message,
    ];
    // your logic here to decide whether to continue the task
    const shouldContinue = Math.random(); /* <-- replace with your logic here */
    enqueueOutput({
      taskStep: step,
      output: response,
      isLast: !shouldContinue,
    });
    if (shouldContinue) {
      const content = await someHeavyFunctionCall();
      // if you want to continue the task, you can insert your new context for the next task step
      step.context.store.messages = [
        ...step.context.store.messages,
        {
          content,
          role: "user",
        },
      ];
    }
  };
}

import { createServer } from "http";
import { OpenAIAgent } from "llamaindex";
import { OpenAIStream, streamToResponse } from "ai";
import { tools } from "./tools";

const agent = new OpenAIAgent({
  tools: [...tools],
});

const server = createServer(async (req, res) => {
  const response = await agent.chat({
    message: "What is weather today?",
    stream: true,
  });

  // Transform the response into a string readable stream
  const stream: ReadableStream<string> = response.pipeThrough(
    new TransformStream({
      transform: (chunk, controller) => {
        controller.enqueue(chunk.response.delta);
      },
    }),
  );

  // Pipe the stream to the response
  streamToResponse(stream, res);
});

server.listen(3000);

type Input = {
  a: number;
  b: number;
};

const sumNumbers = FunctionTool.from<Input>(
  ({ a, b }) => `${a + b}`, // a and b will be checked as number
  // JSON schema will be an error if you type wrong.
  {
    name: "sumNumbers",
    description: "Use this function to sum two numbers",
    parameters: {
      type: "object",
      properties: {
        a: {
          type: "number",
          description: "The first number",
        },
        b: {
          type: "number",
          description: "The second number",
        },
      },
      required: ["a", "b"],
    },
  },
);

// file: src/actions.ts
"use server";
import { Document, VectorStoreIndex } from "llamaindex";
import { readFile } from "node:fs/promises";

const path = "node_modules/llamaindex/examples/abramov.txt";

const essay = await readFile(path, "utf-8");

// Create Document object with essay
const document = new Document({ text: essay, id_: path });

// Split text and create embeddings. Store them in a VectorStoreIndex
const index = await VectorStoreIndex.fromDocuments([document]);

const queryEngine = index.asQueryEngine();

export async function chatWithAI(question: string): Promise<string> {
  const { response } = await queryEngine.query({ query: question });
  return response;
}

export function getWeather(city: string) {
  return `The weather in ${city} is sunny.`;
}

// you don't need to worry about the shcema with different llm tools
export function getTemperature(city: string) {
  return `The temperature in ${city} is 25°C.`;
}

export function getCurrentCity() {
  return "New York";
}

"use server";
import { OpenAI } from "openai";
import { getTools } from "@llamaindex/tool";

export async function chat(message: string) {
  const openai = new OpenAI();
  openai.chat.completions.create({
    messages: [
      {
        role: "user",
        content: "What is the weather in the current city?",
      },
    ],
    tools: getTools("openai"),
  });
}

// next.config.js
const withTool = require("@llamaindex/tool/next");

const config = {
  // Your original Next.js config
};
module.exports = withTool(config);