import { Annotation } from "@langchain/langgraph";

const StateAnnotation = Annotation.Root({
  foo: Annotation<string>,
  bar: Annotation<number>,
});

const StateAnnotation = Annotation.Root({
  messages: Annotation<string[]>({
    reducer: (state, update) => state.concat(update),
    default: () => [],
  }),
});

import { MessagesAnnotation } from "@langchain/langgraph";

const graph = new StateGraph(MessagesAnnotation);

import { Annotation, messagesStateReducer } from "@langchain/langgraph";

const StateAnnotation = Annotation.Root({
  messages: Annotation<BaseMessage[]>({
    reducer: messagesStateReducer,
    default: () => [],
  }),
});

export const GraphState = Annotation.Root({
  messages: Annotation<BaseMessage[]>({
    reducer: (x, y) => x.concat(y),
    default: () => [],
  }),
});

const urls = ["https://policies.google.com/terms?hl=zh-CN"];

const docs = await Promise.all(
  urls.map((url) => new CheerioWebBaseLoader(url).load()),
);
const docsList = docs.flat();

const textSplitter = new RecursiveCharacterTextSplitter({
  chunkSize: 500,
  chunkOverlap: 50,
});
const docSplits = await textSplitter.splitDocuments(docsList);

const vectorStore = await MemoryVectorStore.fromDocuments(
  docSplits,
  new OpenAIEmbeddings({
    configuration: {
      baseURL: "https://api.chatanywhere.tech/v1", //需要更改 Open AI base URL
    },
  }),
);

const retriever = vectorStore.asRetriever();
const tool = createRetrieverTool(retriever, {
  name: "retrieve_blog_posts",
  description:
    "Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.",
});

export const tools = [tool];
export const toolNode = new ToolNode<typeof GraphState.State>(tools);

function shouldRetrieve(state: typeof GraphState.State): string {
  const { messages } = state;
  const lastMessage = messages[messages.length - 1];

  if (
    "tool_calls" in lastMessage &&
    Array.isArray(lastMessage.tool_calls) &&
    lastMessage.tool_calls.length
  ) {
    return "retrieve";
  }
  return END;
}

async function gradeDocuments(state: typeof GraphState.State) {
  const tool = {
    name: "give_relevance_score",
    description: "Give a relevance score to the retrieved documents.",
    schema: z.object({
      binaryScore: z.string().describe("Relevance score 'yes' or 'no'"),
    }),
  };

  const prompt = ChatPromptTemplate.fromTemplate(`
    You are a grader assessing relevance of retrieved docs to a user question.
    ...
    Give a binary score 'yes' or 'no' ...
  `);

  const model = new ChatOpenAI({ ... }).bindTools([tool], {
    tool_choice: tool.name,
  });

  const chain = prompt.pipe(model);
  const lastMessage = state.messages[state.messages.length - 1];

  const score = await chain.invoke({
    question: state.messages[0].content,
    context: lastMessage.content,
  });

  return { messages: [score] };
}

function checkRelevance(state: typeof GraphState.State): string {
  const lastMessage = state.messages[state.messages.length - 1];
  const toolCalls = (lastMessage as AIMessage).tool_calls;
  if (toolCalls[0].args.binaryScore === "yes") {
    return "yes";
  }
  return "no";
}

async function agent(state: typeof GraphState.State) {
  const filteredMessages = state.messages.filter((msg) => {
    if ("tool_calls" in msg) {
      return msg.tool_calls[0].name !== "give_relevance_score";
    }
    return true;
  });

  const model = new ChatOpenAI({ ... }).bindTools(tools);
  const response = await model.invoke(filteredMessages);

  return { messages: [response] };
}

async function rewrite(state: typeof GraphState.State) {
  const question = state.messages[0].content;
  const prompt = ChatPromptTemplate.fromTemplate(`
    Look at the input and try to reason about the underlying semantic intent ...
    Formulate an improved question:
  `);

  const model = new ChatOpenAI({ ... });
  const response = await prompt.pipe(model).invoke({ question });

  return { messages: [response] };
}

async function generate(state: typeof GraphState.State) {
  const question = state.messages[0].content;
  const lastToolMessage = state.messages
    .slice()
    .reverse()
    .find((msg) => msg._getType() === "tool");

  const docs = lastToolMessage.content;
  const prompt = await pull<ChatPromptTemplate>("rlm/rag-prompt");

  const llm = new ChatOpenAI({ ... });
  const ragChain = prompt.pipe(llm);
  const response = await ragChain.invoke({ context: docs, question });

  return { messages: [response] };
}

import { END, StateGraph, START } from "@langchain/langgraph";
import { GraphState } from "./state.js";
import {
  agent,
  gradeDocuments,
  rewrite,
  generate,
  shouldRetrieve,
  checkRelevance,
} from "./edges.js";
import { toolNode } from "./retriever.js";

const workflow = new StateGraph(GraphState)
  .addNode("agent", agent)
  .addNode("retrieve", toolNode)
  .addNode("gradeDocuments", gradeDocuments)
  .addNode("rewrite", rewrite)
  .addNode("generate", generate);

workflow.addEdge(START, "agent");

workflow.addConditionalEdges("agent", shouldRetrieve);

workflow.addEdge("retrieve", "gradeDocuments");

workflow.addEdge("generate", END);
workflow.addEdge("rewrite", "agent");

export const graph = workflow.compile();