feat: rough implementation of langchain prompts

interrobangc · Mar 2, 2024 · a123cde · a123cde
1 parent 2582034
commit a123cde
Show file tree

Hide file tree

Showing 11 changed files with 304 additions and 123 deletions.
diff --git a/core/package.json b/core/package.json
@@ -17,11 +17,13 @@
     "@langchain/core": "^0.1.40",
     "crypto": "^1.0.1",
     "globby": "^14.0.1",
+    "js-yaml": "^4.1.0",
     "lodash": "^4.17.21",
     "winston": "^3.11.0",
     "zod": "^3.22.4"
   },
   "devDependencies": {
+    "@types/js-yaml": "^4.0.9",
     "@types/lodash": "^4.14.202",
     "tsc-alias": "^1.8.8"
   }

diff --git a/core/src/agent/chat.test.ts b/core/src/agent/chat.test.ts
@@ -1,4 +1,5 @@
 import { describe, expect, it, jest } from '@jest/globals';
+import { dump as dumpYaml } from 'js-yaml';
 
 import { decodeResponse } from './chat';
 
@@ -7,10 +8,16 @@ jest.mock('@/log/index.js');
 describe('decodeResponse', () => {
   it('should decode a valid response', () => {
     const response = {
-      type: 'response',
-      content: 'some content',
+      type: 'tool',
+      reason:
+        'To provide context for answering user questions, we can use the `codeSummaryQuery` tool, which searches relevant code snippets and summaries in the vector database. This tool will help us understand the codebase better.',
+      name: 'codeSummaryQuery',
+      params: {
+        includeCode: true,
+        query: 'parseFiles',
+      },
     };
-    const encodedResponse = JSON.stringify(response);
+    const encodedResponse = dumpYaml(response);
     expect(decodeResponse(encodedResponse)).toEqual(response);
   });
 
@@ -28,7 +35,7 @@ describe('decodeResponse', () => {
       text: 'some content',
     };
 
-    const encodedResponse = JSON.stringify(response);
+    const encodedResponse = dumpYaml(response);
     expect(decodeResponse(encodedResponse)).toEqual({
       type: 'error',
       content: encodedResponse,

diff --git a/core/src/agent/chat.ts b/core/src/agent/chat.ts
@@ -1,10 +1,13 @@
 import type { LlmClient, Llms, MessageList, Tools } from '@/.';
 
+import { load as loadYaml } from 'js-yaml';
 import { conversation } from '@/llm/index.js';
 import log from '@/log/index.js';
-import { getPrompt } from '@/prompt/index.js';
+import { newPrompt } from '@/prompt/index.js';
 import * as agentTypes from './types.js';
 
+const prompt = newPrompt();
+
 /**
  * Send a chat message to the LLM
  *
@@ -28,7 +31,7 @@ export const decodeResponse = (
   content: string,
 ): agentTypes.AgentSelectToolResponse => {
   try {
-    return agentTypes.agentLlmResponseSchema.parse(JSON.parse(content.trim()));
+    return agentTypes.agentLlmResponseSchema.parse(loadYaml(content.trim()));
   } catch (e) {
     log('Error decoding response', 'debug', { content, e });
     return {
@@ -38,17 +41,74 @@ export const decodeResponse = (
   }
 };
 
-export const handleMessage = async (
-  llms: Llms,
-  message: string,
-  tools: Tools | undefined,
-  depth: number = 0,
-): Promise<agentTypes.AgentResponse> => {
+export const getToolResponses = (
+  toolResponses: agentTypes.AgentToolResponses,
+) => {
+  return Object.entries(toolResponses)
+    .map(
+      ([tool, response]) => `####${tool}:
+    ${response}
+  `,
+    )
+    .join('\n');
+};
+
+export const handleToolResponse = async ({
+  llms,
+  response,
+  toolResponses,
+  tools,
+}: agentTypes.AgentHandleToolResponseParams): Promise<agentTypes.AgentToolResponses> => {
+  if (!agentTypes.isAgentToolResponse(response)) return toolResponses || {};
+  const toolName = response.name;
+
+  const tool = tools?.[toolName];
+
+  if (!tool) {
+    log('Tool not found', 'error', { toolName, tools });
+    return {
+      ...toolResponses,
+      [toolName]: 'Tool not found',
+    };
+  }
+
+  log(`Running the ${response.name} tool`);
+
+  let toolResponse;
+  try {
+    toolResponse = await tool.run({
+      llm: llms.tool,
+      params: response.params,
+    });
+  } catch (e) {
+    log('Error running tool', 'error', { toolName, e });
+    return {
+      ...toolResponses,
+      [toolName]: 'Error running tool: ' + e,
+    };
+  }
+
+  return {
+    ...toolResponses,
+    [toolName]: toolResponse.content,
+  };
+};
+
+export const handleQuestion = async ({
+  llms,
+  question,
+  toolResponses = {},
+  tools,
+  depth = 0,
+}: agentTypes.AgentHandleQuestionParams): Promise<agentTypes.AgentResponse> => {
   const messages: MessageList = [];
 
   messages.push({
     role: 'user',
-    content: message,
+    content: await prompt.get('agentQuestion', {
+      question,
+      toolResponses: getToolResponses(toolResponses),
+    }),
   });
 
   const response = decodeResponse(await sendChat(llms.agent, messages));
@@ -58,15 +118,22 @@ export const handleMessage = async (
   });
 
   if (agentTypes.isAgentErrorResponse(response)) {
-    conversation.addMessages('agent', [
-      {
-        role: 'user',
-        content: `
-          Be sure to use the correct json format in all further responses.
-        `,
-      },
-    ]);
-    return handleMessage(llms, message, tools, depth + 1);
+    // conversation.addMessages('agent', [
+    //   {
+    //     role: 'user',
+    //     content: `
+    //       Be sure to use the correct yaml format in all further responses.
+    //       ${prompt.get('selectToolFormats')}
+    //     `,
+    //   },
+    // ]);
+    return handleQuestion({
+      llms,
+      question,
+      toolResponses,
+      tools,
+      depth: depth + 1,
+    });
   }
 
   if (agentTypes.isAgentResponseResponse(response)) {
@@ -81,49 +148,18 @@ export const handleMessage = async (
   }
 
   // eslint-disable-next-line  @typescript-eslint/no-use-before-define
-  return handleToolResponse(llms, message, response, tools, depth);
-};
-
-export const handleToolResponse = async (
-  llms: Llms,
-  message: string,
-  response: agentTypes.AgentSelectToolResponse,
-  tools: Tools | undefined,
-  depth: number = 0,
-) => {
-  if (!agentTypes.isAgentToolResponse(response)) return response;
-
-  const tool = tools?.[response.name];
-
-  if (!tool) {
-    log('Tool not found', 'debug', { response, tools });
-    const toolNotFoundMessage = `
-      ${message}
-
-      The ${response.name} tool was not found. Please select a different tool.
-    `;
-
-    return handleMessage(llms, toolNotFoundMessage, tools, depth + 1);
-  }
-
-  log(`Running the ${response.name} tool`);
-
-  const toolResponse = await tool.run({
-    llm: llms.tool,
-    params: response.params,
+  return handleQuestion({
+    llms,
+    question,
+    toolResponses: await handleToolResponse({
+      llms,
+      response,
+      toolResponses,
+      tools,
+    }),
+    tools,
+    depth: depth + 1,
   });
-
-  const toolResponseMessage = `
-    ${message}
-    ${getPrompt('toolResponseStart')}
-
-    Tool: ${response.name}
-
-    Response:
-    ${toolResponse.content}
-  `;
-
-  return handleMessage(llms, toolResponseMessage, tools, depth + 1);
 };
 
 /**
@@ -139,12 +175,12 @@ export const handleToolResponse = async (
 export const chat =
   (llms: Llms, tools: Tools | undefined) =>
   async (question: string): Promise<agentTypes.AgentResponse> => {
-    const message = `
-      ${getPrompt('userQuestionStart')}
-      ${question}
-    `;
-
-    return handleMessage(llms, message, tools);
+    log('chat', 'debug', { question });
+    return handleQuestion({
+      llms,
+      question,
+      tools,
+    });
   };
 
 export default chat;
diff --git a/core/src/agent/newAgent.ts b/core/src/agent/newAgent.ts
@@ -1,15 +1,11 @@
-import type { Agent, PartialConfig, Tools } from '@/.';
+import type { Agent, PartialConfig } from '@/.';
 import { getConfig, initConfig } from '@/config/index.js';
 import { conversation, initLlms } from '@/llm/index.js';
 import { log } from '@/log/index.js';
-import { getPrompt } from '@/prompt/index.js';
+import { initPrompts } from '@/prompt/index.js';
 import { initTools } from '@/tool/index.js';
 import chat from './chat.js';
 
-export const getToolDescriptions = (tools: Tools = {}) => {
-  return JSON.stringify(Object.values(tools).map((tool) => tool.description));
-};
-
 /**
  * Create a new agent which is the primary interface to interact with the LLMs
  *
@@ -27,14 +23,12 @@ export const newAgent = async (configParam: PartialConfig): Promise<Agent> => {
   const tools = await initTools(config);
   log('newAgent tools', 'silly', { tools });
 
+  const prompt = initPrompts({ tools });
+
   conversation.addMessages('agent', [
     {
       role: 'system',
-      content: `
-        ${getPrompt('agent')}
-        ${getPrompt('selectTool')}
-        ${getToolDescriptions(tools)}
-      `,
+      content: await prompt.get('agentSystem'),
     },
   ]);
 

diff --git a/core/src/agent/types.ts b/core/src/agent/types.ts
@@ -1,3 +1,5 @@
+import type { Llms, Tools } from '@/.';
+
 import { z } from 'zod';
 
 import { toolRunParamsParamSchema } from '@/tool/types.js';
@@ -28,7 +30,6 @@ export const agentToolResponseSchema = z.object({
   type: z.literal('tool'),
   reason: z.string(),
   name: z.string(),
-  query: z.string(),
   params: toolRunParamsParamSchema,
 });
 
@@ -73,3 +74,20 @@ export type AgentSelectToolResponse =
   | AgentToolResponse;
 
 export type AgentResponse = AgentErrorResponse | AgentResponseResponse;
+
+export type AgentToolResponses = Record<string, string>;
+
+export type AgentHandleQuestionParams = {
+  llms: Llms;
+  question: string;
+  toolResponses?: AgentToolResponses;
+  tools: Tools | undefined;
+  depth?: number;
+};
+
+export type AgentHandleToolResponseParams = {
+  llms: Llms;
+  response: AgentSelectToolResponse;
+  toolResponses: AgentToolResponses;
+  tools: Tools | undefined;
+};
diff --git a/core/src/index.ts b/core/src/index.ts
@@ -19,5 +19,6 @@ export * from './config/types.js';
 export * from './importer/types.js';
 export * from './llm/types.js';
 export * from './log/types.js';
+export * from './prompt/types.js';
 export * from './tool/types.js';
 export * from './vectorDb/types.js';