agentlang 0.0.2

package/src/runtime/modules/ai.ts

@@ -0,0 +1,257 @@
+import { makeCoreModuleName, makeFqName } from '../util.js';
+import { Environment, makeEventEvaluator, parseAndEvaluateStatement } from '../interpreter.js';
+import { fetchModule, Instance, instanceToObject, isModule } from '../module.js';
+import { provider } from '../agents/registry.js';
+import {
+  AgentServiceProvider,
+  AIResponse,
+  assistantMessage,
+  humanMessage,
+  systemMessage,
+} from '../agents/provider.js';
+import { AIMessage, BaseMessage, HumanMessage } from '@langchain/core/messages';
+import { PlannerInstructions } from '../agents/common.js';
+import { PathAttributeNameQuery } from '../defs.js';
+
+export const CoreAIModuleName = makeCoreModuleName('ai');
+
+export default `module ${CoreAIModuleName}
+
+entity llm {
+    name String @id,
+    service String @default("openai"),
+    config Map @optional
+}
+
+entity agent {
+    name String @id,
+    type @enum("chat", "planner") @default("chat"),
+    instruction String @optional,
+    tools String[] @optional,
+    documents String[] @optional,
+    llm String
+}
+
+entity agentChatSession {
+    id String @id,
+    messages String
+}
+
+workflow findAgentChatSession {
+  {agentChatSession {id? findAgentChatSession.id}} as [sess];
+  sess
+}
+
+workflow saveAgentChatSession {
+  upsert {agentChatSession {id saveAgentChatSession.id, messages saveAgentChatSession.messages}}
+}
+
+entity document {
+  title String @id,
+  content String,
+  @meta {"fullTextSearch": "*"}
+}
+`;
+
+export const AgentFqName = makeFqName(CoreAIModuleName, 'agent');
+
+const ProviderDb = new Map<string, AgentServiceProvider>();
+
+export class Agent {
+  llm: string = '';
+  name: string = '';
+  chatId: string | undefined;
+  instruction: string = '';
+  type: string = 'chat';
+  tools: string[] | undefined;
+  documents: string[] | undefined;
+
+  private constructor() {}
+
+  static FromInstance(agentInstance: Instance): Agent {
+    return instanceToObject<Agent>(agentInstance, new Agent());
+  }
+
+  isPlanner(): boolean {
+    return (this.tools && this.tools.length > 0) || this.type == 'planner';
+  }
+
+  async invoke(message: string, env: Environment) {
+    const p = await findProviderForLLM(this.llm, env);
+    const agentName = this.name;
+    const chatId = this.chatId || agentName;
+    const sess: Instance | null = await findAgentChatSession(chatId, env);
+    let msgs: BaseMessage[] | undefined;
+    const isplnr = this.isPlanner();
+    if (sess) {
+      msgs = sess.lookup('messages');
+    } else {
+      msgs = [systemMessage(this.instruction)];
+    }
+    if (msgs) {
+      const sysMsg = msgs[0];
+      if (isplnr) {
+        const newSysMsg = systemMessage(
+          `${PlannerInstructions}\n${this.toolsAsString()}\n${this.instruction}`
+        );
+        msgs[0] = newSysMsg;
+      }
+      msgs.push(humanMessage(await this.maybeAddRelevantDocuments(message, env)));
+      const response: AIResponse = await p.invoke(msgs);
+      msgs.push(assistantMessage(response.content));
+      if (isplnr) {
+        msgs[0] = sysMsg;
+      }
+      await saveAgentChatSession(chatId, msgs, env);
+      env.setLastResult(response.content);
+    } else {
+      throw new Error(`failed to initialize messages for agent ${agentName}`);
+    }
+  }
+
+  private async maybeAddRelevantDocuments(message: string, env: Environment): Promise<string> {
+    if (this.documents && this.documents.length > 0) {
+      const s = `${message}. Relevant documents are: ${this.documents.join(',')}`;
+      const result: any[] = await parseHelper(`{agentlang_ai/document? "${s}"}`, env);
+      if (result && result.length > 0) {
+        const docs: Instance[] = [];
+        for (let i = 0; i < result.length; ++i) {
+          const v: any = result[i];
+          const r: Instance[] = await parseHelper(
+            `{agentlang_ai/document {${PathAttributeNameQuery} "${v.id}"}}`,
+            env
+          );
+          if (r && r.length > 0) {
+            docs.push(r[0]);
+          }
+        }
+        if (docs.length > 0) {
+          message = message.concat('\nUse the additional information given below:\n').concat(
+            docs
+              .map((v: Instance) => {
+                return v.lookup('content');
+              })
+              .join('\n')
+          );
+        }
+      }
+    }
+    return message;
+  }
+
+  private toolsAsString(): string {
+    if (this.tools) {
+      return this.tools
+        .filter((s: string) => {
+          return isModule(s);
+        })
+        .map((moduleName: string) => {
+          return fetchModule(moduleName).toString();
+        })
+        .join('\n');
+    } else {
+      return '';
+    }
+  }
+}
+
+async function parseHelper(stmt: string, env: Environment): Promise<any> {
+  await parseAndEvaluateStatement(stmt, undefined, env);
+  return env.getLastResult();
+}
+
+export async function findAgentByName(name: string, env: Environment): Promise<Agent> {
+  const result = await parseHelper(`{agentlang_ai/agent {name? "${name}"}}`, env);
+  if (result instanceof Array && result.length > 0) {
+    const agentInstance: Instance = result[0];
+    return Agent.FromInstance(agentInstance);
+  } else {
+    throw new Error(`Failed to fine agent ${name}`);
+  }
+}
+
+export async function findProviderForLLM(
+  llmName: string,
+  env: Environment
+): Promise<AgentServiceProvider> {
+  let p: AgentServiceProvider | undefined = ProviderDb.get(llmName);
+  if (p == undefined) {
+    const result: Instance[] = await parseAndEvaluateStatement(
+      `{${CoreAIModuleName}/llm {name? "${llmName}"}}`,
+      undefined,
+      env
+    );
+    if (result.length > 0) {
+      const llm: Instance = result[0];
+      const service = llm.lookup('service');
+      const pclass = provider(service);
+      const providerConfig: Map<string, any> =
+        llm.lookup('config') || new Map().set('service', service);
+      p = new pclass(providerConfig);
+      if (p) ProviderDb.set(llmName, p);
+    }
+  }
+  if (p) {
+    return p;
+  } else {
+    throw new Error(`Failed to load provider for ${llmName}`);
+  }
+}
+
+const evalEvent = makeEventEvaluator(CoreAIModuleName);
+
+type GenericMessage = {
+  role: 'system' | 'user' | 'assistant';
+  content: string;
+};
+
+function asBaseMessages(gms: GenericMessage[]): BaseMessage[] {
+  return gms.map((gm: GenericMessage): BaseMessage => {
+    switch (gm.role) {
+      case 'user': {
+        return humanMessage(gm.content);
+      }
+      case 'assistant': {
+        return assistantMessage(gm.content);
+      }
+      default: {
+        return systemMessage(gm.content);
+      }
+    }
+  });
+}
+
+function asGenericMessages(bms: BaseMessage[]): GenericMessage[] {
+  return bms.map((bm: BaseMessage): GenericMessage => {
+    if (bm instanceof HumanMessage) {
+      return { role: 'user', content: bm.text };
+    } else if (bm instanceof AIMessage) {
+      return { role: 'assistant', content: bm.text };
+    } else {
+      return { role: 'system', content: bm.text };
+    }
+  });
+}
+
+export async function findAgentChatSession(
+  chatId: string,
+  env: Environment
+): Promise<Instance | null> {
+  const result: Instance | null = await evalEvent('findAgentChatSession', { id: chatId }, env);
+  if (result) {
+    result.attributes.set('messages', asBaseMessages(JSON.parse(result.lookup('messages'))));
+  }
+  return result;
+}
+
+export async function saveAgentChatSession(chatId: string, messages: any[], env: Environment) {
+  await evalEvent(
+    'saveAgentChatSession',
+    { id: chatId, messages: JSON.stringify(asGenericMessages(messages)) },
+    env
+  );
+}
+
+export function agentName(agentInstance: Instance): string {
+  return agentInstance.lookup('name');
+}