From b495a69d6ad62c8b8f713baae10e34323bfce030 Mon Sep 17 00:00:00 2001 From: Stephen Chen Date: Tue, 22 Apr 2025 00:52:17 +0800 Subject: [PATCH] Prototype feature: allow agents to go back-and-forth with AI chat in a thread --- .gitignore | 2 ++ package-lock.json | 42 ++++++++++++++++++++++++++++++++++++++++++ package.json | 7 +++++-- src/api.test.ts | 4 +++- src/api.ts | 9 +++++++-- src/db.ts | 20 ++++++++++++++++++++ src/index.ts | 6 ++++-- src/mcp.ts | 47 ++++++++++++++++++++++++++++++++++++++++------- 8 files changed, 123 insertions(+), 14 deletions(-) create mode 100644 src/db.ts diff --git a/.gitignore b/.gitignore index 5432641..9dfb946 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ dist/ node_modules/ +.env +db.json # System .DS_Store diff --git a/package-lock.json b/package-lock.json index 7a0f182..0e16c80 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,6 +9,8 @@ "version": "0.1.0", "dependencies": { "@modelcontextprotocol/sdk": "^1.8.0", + "lowdb": "^7.0.1", + "uuid": "^11.1.0", "zod": "^3.24.2" }, "bin": { @@ -4606,6 +4608,21 @@ "dev": true, "license": "MIT" }, + "node_modules/lowdb": { + "version": "7.0.1", + "resolved": "https://registry.npmjs.org/lowdb/-/lowdb-7.0.1.tgz", + "integrity": "sha512-neJAj8GwF0e8EpycYIDFqEPcx9Qz4GUho20jWFR7YiFeXzF1YMLdxB36PypcTSPMA+4+LvgyMacYhlr18Zlymw==", + "license": "MIT", + "dependencies": { + "steno": "^4.0.2" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "url": "https://github.com/sponsors/typicode" + } + }, "node_modules/lru-cache": { "version": "10.4.3", "resolved": "https://registry.npmjs.org/lru-cache/-/lru-cache-10.4.3.tgz", @@ -5751,6 +5768,18 @@ "node": ">= 0.8" } }, + "node_modules/steno": { + "version": "4.0.2", + "resolved": "https://registry.npmjs.org/steno/-/steno-4.0.2.tgz", + "integrity": "sha512-yhPIQXjrlt1xv7dyPQg2P17URmXbuM5pdGkpiMB3RenprfiBlvK415Lctfe0eshk90oA7/tNq7WEiMK8RSP39A==", + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "url": "https://github.com/sponsors/typicode" + } + }, "node_modules/string-length": { "version": "4.0.2", "resolved": "https://registry.npmjs.org/string-length/-/string-length-4.0.2.tgz", @@ -6412,6 +6441,19 @@ "browserslist": ">= 4.21.0" } }, + "node_modules/uuid": { + "version": "11.1.0", + "resolved": "https://registry.npmjs.org/uuid/-/uuid-11.1.0.tgz", + "integrity": "sha512-0/A9rDy9P7cJ+8w1c9WD9V//9Wj15Ce2MPz8Ri6032usz+NfePxx5AcN3bN+r6ZL6jEo066/yNYB3tn4pQEx+A==", + "funding": [ + "https://github.com/sponsors/broofa", + "https://github.com/sponsors/ctavan" + ], + "license": "MIT", + "bin": { + "uuid": "dist/esm/bin/uuid" + } + }, "node_modules/v8-to-istanbul": { "version": "9.3.0", "resolved": "https://registry.npmjs.org/v8-to-istanbul/-/v8-to-istanbul-9.3.0.tgz", diff --git a/package.json b/package.json index 2f22982..e0b4dea 100644 --- a/package.json +++ b/package.json @@ -2,6 +2,7 @@ "name": "fern-mcp-server", "description": "Model Context Protocol (MCP) server for the Fern API.", "version": "0.1.0", + "type": "module", "bin": "dist/index.js", "files": [ "dist" @@ -9,12 +10,14 @@ "scripts": { "start": "concurrently \"npm run build:watch\" \"npm run inspector\"", "inspector": "npx @modelcontextprotocol/inspector@latest -- nodemon --env-file=.env -q --watch dist dist/index.js", - "build": "tsup src/index.ts --dts --clean", - "build:watch": "tsup src/index.ts --dts --watch", + "build": "tsup src/index.ts --format esm --dts --clean", + "build:watch": "tsup src/index.ts --format esm --dts --watch", "test": "jest" }, "dependencies": { "@modelcontextprotocol/sdk": "^1.8.0", + "lowdb": "^7.0.1", + "uuid": "^11.1.0", "zod": "^3.24.2" }, "devDependencies": { diff --git a/src/api.test.ts b/src/api.test.ts index 3199b07..275c0b5 100644 --- a/src/api.test.ts +++ b/src/api.test.ts @@ -6,7 +6,9 @@ describe("postChat", () => { it( "returns a non-empty string", async () => { - const result = await postChat("What is Fern AI Chat?"); + const result = await postChat([ + { role: "user", content: "What is Fern AI Chat?" }, + ]); console.log({ result }); expect(result).not.toBe(""); diff --git a/src/api.ts b/src/api.ts index d44f439..f325fda 100644 --- a/src/api.ts +++ b/src/api.ts @@ -1,13 +1,18 @@ const BASE_URL = "https://buildwithfern.com/learn"; -export async function postChat(message: string) { +interface ChatMessage { + role: "user" | "assistant"; + content: string; +} + +export async function postChat(messages: ChatMessage[]) { const response = await fetch(`${BASE_URL}/api/fern-docs/search/v2/chat`, { method: "POST", headers: { "content-type": "application/json", }, body: JSON.stringify({ - messages: [{ role: "user", content: message }], + messages: messages, url: "https://buildwithfern.com/learn", filters: [], }), diff --git a/src/db.ts b/src/db.ts new file mode 100644 index 0000000..47f5921 --- /dev/null +++ b/src/db.ts @@ -0,0 +1,20 @@ +import { Low, Memory } from "lowdb"; + +export type Thread = Array<{ + role: "user" | "assistant"; + content: string; +}>; + +export type Data = { + threads: Record; +}; + +export type DB = Low; + +const defaultData: Data = { threads: {} }; + +export function createDb() { + // TODO: Use JSONFilePreset instead + const db = new Low(new Memory(), defaultData); + return db; +} diff --git a/src/index.ts b/src/index.ts index 68135b1..5195595 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,7 +1,8 @@ #!/usr/bin/env node import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; -import { createMcpServer, registerMcpTools } from "./mcp"; +import { createMcpServer, registerMcpTools } from "./mcp.js"; +import { createDb } from "./db.js"; const packageJson = require("../package.json") as any; @@ -12,7 +13,8 @@ async function run() { } const server = createMcpServer(packageJson.name, packageJson.version); - registerMcpTools(server); + const db = createDb(); + registerMcpTools(server, db); const transport = new StdioServerTransport(); await server.connect(transport); diff --git a/src/mcp.ts b/src/mcp.ts index c4bc4d1..91ab31a 100644 --- a/src/mcp.ts +++ b/src/mcp.ts @@ -1,8 +1,10 @@ #!/usr/bin/env node import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { v4 as uuidv4 } from "uuid"; import { z } from "zod"; -import * as api from "./api"; +import * as api from "./api.js"; +import { DB } from "./db.js"; // Create an MCP server export function createMcpServer(name: string, version: string) { @@ -19,15 +21,46 @@ ABOUT FERN (builtwithfern.com): Start with OpenAPI. Generate SDKs in multiple la } // Register MCP tools -export function registerMcpTools(server: McpServer) { +export function registerMcpTools(server: McpServer, db: DB) { server.tool( "ask_fern_ai", - "Ask Fern AI about anything related to Fern.", - { message: z.string() }, - async ({ message }) => { - const result = await api.postChat(message); + `Ask Fern AI about anything related to Fern. +Don't include a threadId in your initial message. +If a message requires a follow-up, include the threadId in your follow-up messages until the thread is resolved.`, + { question: z.string(), threadId: z.string().optional() }, + async ({ question, threadId: _threadId }) => { + // Issue a new thread ID if none is provided + const threadId = _threadId ?? uuidv4(); + await db.read(); + const thread = [...(db.data.threads[threadId] ?? [])]; + + // Create a user message + thread.push({ role: "user" as const, content: question }); + + // Post the question to the API + const answer = await api.postChat(thread); + + // Create an assistant message + thread.push({ role: "assistant" as const, content: answer }); + + console.error(thread); + + // Update the thread state + db.update(({ threads }) => { + threads[threadId] = thread; + }); + + // Return the answer (last message in the thread) return { - content: [{ type: "text", text: result }], + content: [ + { + type: "text", + text: JSON.stringify({ + threadId: threadId, + thread: thread, + }), + }, + ], }; } );