feat(agents): Support LLMChain in OpenAIFunctionsAgent and memory

davidmigloz · Aug 5, 2023 · c5b0a0a · c5b0a0a
1 parent edf4ea5
commit c5b0a0a
Show file tree

Hide file tree

Showing 13 changed files with 256 additions and 93 deletions.
diff --git a/packages/langchain/lib/src/agents/agent.dart b/packages/langchain/lib/src/agents/agent.dart
@@ -9,6 +9,9 @@ abstract class BaseActionAgent {
   /// {@macro base_action_agent}
   const BaseActionAgent();
 
+  /// The key for the scratchpad (intermediate steps) of the agent.
+  static const agentScratchpadInputKey = 'agent_scratchpad';
+
   /// Return key for the agent's output.
   static const agentReturnKey = 'output';
 

diff --git a/packages/langchain/lib/src/agents/executors.dart b/packages/langchain/lib/src/agents/executors.dart
@@ -17,6 +17,12 @@ import 'tools/invalid.dart';
 /// retrieves the output, and passes it back to the agent to determine the next
 /// action. This process continues until the agent determines it can directly
 /// respond to the user or completes its task.
+///
+/// If you add [memory] to the [AgentExecutor], it will save the
+/// [AgentExecutor]'s inputs and outputs. It won't save the agent's
+/// intermediate inputs and outputs. If you want to save the agent's
+/// intermediate inputs and outputs, you should add [memory] to the agent
+/// instead.
 /// {@endtemplate}
 class AgentExecutor extends BaseChain {
   AgentExecutor({

diff --git a/packages/langchain/lib/src/agents/tools/base.dart b/packages/langchain/lib/src/agents/tools/base.dart
@@ -2,6 +2,7 @@ import 'dart:async';
 
 import 'package:meta/meta.dart';
 
+import '../../model_io/chat_models/models/models.dart';
 import 'models/models.dart';
 
 /// {@template base_tool}
@@ -104,6 +105,15 @@ abstract base class BaseTool {
     return run(toolInput);
   }
 
+  /// Converts the tool to a [ChatFunction].
+  ChatFunction toChatFunction() {
+    return ChatFunction(
+      name: name,
+      description: description,
+      parameters: inputJsonSchema,
+    );
+  }
+
   @override
   bool operator ==(covariant final BaseTool other) =>
       identical(this, other) || name == other.name;

diff --git a/packages/langchain/lib/src/memory/chat.dart b/packages/langchain/lib/src/memory/chat.dart
@@ -1,5 +1,6 @@
 import 'dart:async';
 
+import '../model_io/chat_models/models/models.dart';
 import '../utils/exception.dart';
 import 'base.dart';
 import 'models/models.dart';
@@ -46,11 +47,21 @@ abstract base class BaseChatMemory implements BaseMemory {
   }) async {
     // this is purposefully done in sequence so they're saved in order
     final (input, output) = _getInputOutputValues(inputValues, outputValues);
-    await chatHistory.addUserChatMessage(input);
-    await chatHistory.addAIChatMessage(output);
+
+    if (input is ChatMessage) {
+      await chatHistory.addChatMessage(input);
+    } else {
+      await chatHistory.addHumanChatMessage(input.toString());
+    }
+
+    if (output is ChatMessage) {
+      await chatHistory.addChatMessage(output);
+    } else {
+      await chatHistory.addAIChatMessage(output.toString());
+    }
   }
 
-  (String input, String output) _getInputOutputValues(
+  (dynamic input, dynamic output) _getInputOutputValues(
     final MemoryInputValues inputValues,
     final MemoryOutputValues outputValues,
   ) {

diff --git a/packages/langchain/lib/src/memory/stores/message/history.dart b/packages/langchain/lib/src/memory/stores/message/history.dart
@@ -15,8 +15,8 @@ abstract base class BaseChatMessageHistory {
   /// Add [ChatMessage] to the history.
   Future<void> addChatMessage(final ChatMessage message);
 
-  /// Add a user message to the history.
-  Future<void> addUserChatMessage(final String message) {
+  /// Add a human message to the history.
+  Future<void> addHumanChatMessage(final String message) {
     return addChatMessage(ChatMessage.human(message));
   }
 

diff --git a/packages/langchain/lib/src/memory/utils.dart b/packages/langchain/lib/src/memory/utils.dart
@@ -1,3 +1,4 @@
+import '../agents/agent.dart';
 import '../utils/exception.dart';
 import 'models/models.dart';
 
@@ -10,13 +11,17 @@ String getPromptInputKey(
   final MemoryInputValues inputValues,
   final Set<String> memoryKeys,
 ) {
-  // "stop" is a special key that can be passed as input but is not used to
-  // format the prompt
-  final promptInputKeys =
-      inputValues.keys.toSet().difference({...memoryKeys, 'stop'});
+  // Reserved keys can be passed as input but is not used to format the prompt
+  final promptInputKeys = inputValues.keys.toSet().difference({
+    ...memoryKeys,
+    'stop',
+    BaseActionAgent.agentScratchpadInputKey,
+  });
   if (promptInputKeys.length != 1) {
     throw LangChainException(
-      message: 'One input key expected got $promptInputKeys',
+      message: 'One input key expected got $promptInputKeys. '
+          'If you have multiple input keys in your prompt you need to specify '
+          'the input key to use for the memory using the `inputKey` parameter.',
     );
   }
   return promptInputKeys.first;

diff --git a/packages/langchain/lib/src/model_io/chat_models/models/models.dart b/packages/langchain/lib/src/model_io/chat_models/models/models.dart
@@ -336,12 +336,12 @@ CustomChatMessage{
 /// Role of a chat message
 enum ChatMessageRole { system, human, ai, custom }
 
-/// {@template openai_function_model}
+/// {@template chat_function}
 /// The description of a function that can be called by the chat model.
 /// {@endtemplate
 @immutable
 class ChatFunction {
-  /// {@macro openai_function_model}
+  /// {@macro chat_function}
   const ChatFunction({
     required this.name,
     this.description,

diff --git a/packages/langchain/test/memory/buffer_test.dart b/packages/langchain/test/memory/buffer_test.dart
@@ -1,3 +1,4 @@
+import 'package:langchain/src/agents/agent.dart';
 import 'package:langchain/src/memory/memory.dart';
 import 'package:langchain/src/model_io/chat_models/chat_models.dart';
 import 'package:test/test.dart';
@@ -7,21 +8,21 @@ void main() {
     test('Test buffer memory', () async {
       final memory = ConversationBufferMemory();
       final result1 = await memory.loadMemoryVariables();
-      expect(result1, {'history': ''});
+      expect(result1, {BaseMemory.defaultMemoryKey: ''});
 
       await memory.saveContext(
         inputValues: {'foo': 'bar'},
         outputValues: {'bar': 'foo'},
       );
       const expectedString = 'Human: bar\nAI: foo';
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': expectedString});
+      expect(result2, {BaseMemory.defaultMemoryKey: expectedString});
     });
 
     test('Test buffer memory return messages', () async {
       final memory = ConversationBufferMemory(returnMessages: true);
       final result1 = await memory.loadMemoryVariables();
-      expect(result1, {'history': <ChatMessage>[]});
+      expect(result1, {BaseMemory.defaultMemoryKey: <ChatMessage>[]});
 
       await memory.saveContext(
         inputValues: {'foo': 'bar'},
@@ -32,7 +33,24 @@ void main() {
         ChatMessage.ai('foo'),
       ];
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': expectedResult});
+      expect(result2, {BaseMemory.defaultMemoryKey: expectedResult});
+    });
+
+    test('Test chat message as input and output', () async {
+      final memory = ConversationBufferMemory(returnMessages: true);
+      final result1 = await memory.loadMemoryVariables();
+      expect(result1, {BaseMemory.defaultMemoryKey: <ChatMessage>[]});
+
+      await memory.saveContext(
+        inputValues: {'foo': ChatMessage.function(name: 'foo', content: 'bar')},
+        outputValues: {'bar': ChatMessage.ai('baz')},
+      );
+      final expectedResult = [
+        ChatMessage.function(name: 'foo', content: 'bar'),
+        ChatMessage.ai('baz'),
+      ];
+      final result2 = await memory.loadMemoryVariables();
+      expect(result2, {BaseMemory.defaultMemoryKey: expectedResult});
     });
 
     test('Test buffer memory with pre-loaded history', () async {
@@ -45,7 +63,7 @@ void main() {
         chatHistory: ChatMessageHistory(messages: pastMessages),
       );
       final result = await memory.loadMemoryVariables();
-      expect(result, {'history': pastMessages});
+      expect(result, {BaseMemory.defaultMemoryKey: pastMessages});
     });
 
     test('Test clear memory', () async {
@@ -56,11 +74,69 @@ void main() {
       );
       const expectedString = 'Human: bar\nAI: foo';
       final result1 = await memory.loadMemoryVariables();
-      expect(result1, {'history': expectedString});
+      expect(result1, {BaseMemory.defaultMemoryKey: expectedString});
 
       memory.clear();
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': ''});
+      expect(result2, {BaseMemory.defaultMemoryKey: ''});
+    });
+
+    test('Test reserved keys are ignored when selecting prompt input keys',
+        () async {
+      final memory = ConversationBufferMemory(returnMessages: true);
+      await memory.saveContext(
+        inputValues: {
+          'foo': 'bar',
+          'stop': 'stop',
+          BaseActionAgent.agentScratchpadInputKey: 'baz',
+        },
+        outputValues: {'bar': 'foo'},
+      );
+      final expectedResult = [
+        ChatMessage.human('bar'),
+        ChatMessage.ai('foo'),
+      ];
+      final result1 = await memory.loadMemoryVariables();
+      expect(result1, {BaseMemory.defaultMemoryKey: expectedResult});
+    });
+
+    test('Test multiple input values with inputKey specified', () async {
+      final memory = ConversationBufferMemory(
+        returnMessages: true,
+        inputKey: 'foo2',
+      );
+      await memory.saveContext(
+        inputValues: {
+          'foo1': 'bar1',
+          'foo2': 'bar2',
+          BaseActionAgent.agentScratchpadInputKey: 'baz',
+        },
+        outputValues: {'bar': 'foo'},
+      );
+      final expectedResult = [
+        ChatMessage.human('bar2'),
+        ChatMessage.ai('foo'),
+      ];
+      final result1 = await memory.loadMemoryVariables();
+      expect(result1, {BaseMemory.defaultMemoryKey: expectedResult});
+    });
+
+    test(
+        'Test error is thrown if inputKey not specified when using with '
+        'multiple input values', () async {
+      final memory = ConversationBufferMemory(returnMessages: true);
+
+      // expect throws exception if no input keys are selected
+      expect(
+        () async => memory.saveContext(
+          inputValues: {
+            'foo1': 'bar1',
+            'foo2': 'bar2',
+          },
+          outputValues: {'bar': 'foo'},
+        ),
+        throwsException,
+      );
     });
   });
 }
diff --git a/packages/langchain/test/memory/buffer_window_test.dart b/packages/langchain/test/memory/buffer_window_test.dart
@@ -15,13 +15,13 @@ void main() {
       );
       const expectedString = 'Human: bar\nAI: foo';
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': expectedString});
+      expect(result2, {BaseMemory.defaultMemoryKey: expectedString});
     });
 
     test('Test buffer memory return messages', () async {
       final memory = ConversationBufferWindowMemory(k: 1, returnMessages: true);
       final result1 = await memory.loadMemoryVariables();
-      expect(result1, {'history': <ChatMessage>[]});
+      expect(result1, {BaseMemory.defaultMemoryKey: <ChatMessage>[]});
 
       await memory.saveContext(
         inputValues: {'foo': 'bar'},
@@ -32,7 +32,7 @@ void main() {
         ChatMessage.ai('foo'),
       ];
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': expectedResult});
+      expect(result2, {BaseMemory.defaultMemoryKey: expectedResult});
 
       await memory.saveContext(
         inputValues: {'foo': 'bar1'},
@@ -44,7 +44,7 @@ void main() {
         ChatMessage.ai('foo1'),
       ];
       final result3 = await memory.loadMemoryVariables();
-      expect(result3, {'history': expectedResult2});
+      expect(result3, {BaseMemory.defaultMemoryKey: expectedResult2});
     });
 
     test('Test buffer memory with pre-loaded history', () async {
@@ -57,7 +57,7 @@ void main() {
         chatHistory: ChatMessageHistory(messages: pastMessages),
       );
       final result = await memory.loadMemoryVariables();
-      expect(result, {'history': pastMessages});
+      expect(result, {BaseMemory.defaultMemoryKey: pastMessages});
     });
 
     test('Test clear memory', () async {
@@ -68,11 +68,11 @@ void main() {
       );
       const expectedString = 'Human: bar\nAI: foo';
       final result1 = await memory.loadMemoryVariables();
-      expect(result1, {'history': expectedString});
+      expect(result1, {BaseMemory.defaultMemoryKey: expectedString});
 
       memory.clear();
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': ''});
+      expect(result2, {BaseMemory.defaultMemoryKey: ''});
     });
   });
 }
diff --git a/packages/langchain/test/memory/stores/message/in_memory_test.dart b/packages/langchain/test/memory/stores/message/in_memory_test.dart
@@ -13,7 +13,7 @@ void main() {
 
     test('Test addUserMessage', () async {
       final history = ChatMessageHistory()
-        ..addUserChatMessage('This is a human msg');
+        ..addHumanChatMessage('This is a human msg');
       final messages = await history.getChatMessages();
       expect(messages.first, isA<HumanChatMessage>());
       expect(messages.first.content, 'This is a human msg');

diff --git a/packages/langchain/test/memory/token_buffer_test.dart b/packages/langchain/test/memory/token_buffer_test.dart
@@ -18,7 +18,7 @@ void main() {
       );
       const expectedString = 'Human: bar\nAI: foo';
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': expectedString});
+      expect(result2, {BaseMemory.defaultMemoryKey: expectedString});
     });
 
     test('Test buffer memory return messages', () async {
@@ -29,7 +29,7 @@ void main() {
         maxTokenLimit: 4,
       );
       final result1 = await memory.loadMemoryVariables();
-      expect(result1, {'history': <ChatMessage>[]});
+      expect(result1, {BaseMemory.defaultMemoryKey: <ChatMessage>[]});
 
       await memory.saveContext(
         inputValues: {'foo': 'bar'},
@@ -40,7 +40,7 @@ void main() {
         ChatMessage.ai('foo'),
       ];
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': expectedResult});
+      expect(result2, {BaseMemory.defaultMemoryKey: expectedResult});
 
       await memory.saveContext(
         inputValues: {'foo': 'bar1'},
@@ -53,7 +53,7 @@ void main() {
         ChatMessage.ai('foo1'),
       ];
       final result3 = await memory.loadMemoryVariables();
-      expect(result3, {'history': expectedResult2});
+      expect(result3, {BaseMemory.defaultMemoryKey: expectedResult2});
     });
 
     test('Test buffer memory with pre-loaded history', () async {
@@ -69,7 +69,7 @@ void main() {
         chatHistory: ChatMessageHistory(messages: pastMessages),
       );
       final result = await memory.loadMemoryVariables();
-      expect(result, {'history': pastMessages});
+      expect(result, {BaseMemory.defaultMemoryKey: pastMessages});
     });
 
     test('Test clear memory', () async {
@@ -80,11 +80,11 @@ void main() {
       );
       const expectedString = 'Human: bar\nAI: foo';
       final result1 = await memory.loadMemoryVariables();
-      expect(result1, {'history': expectedString});
+      expect(result1, {BaseMemory.defaultMemoryKey: expectedString});
 
       memory.clear();
       final result2 = await memory.loadMemoryVariables();
-      expect(result2, {'history': ''});
+      expect(result2, {BaseMemory.defaultMemoryKey: ''});
     });
   });
 }