Merge pull request shroominic#79 from preritdas/custom-system-message

Allow custom system message, custom max_iterations
tjroamer · Aug 20, 2023 · f6e72e6 · f6e72e6
2 parents 65d48aa + 24ffc9d
commit f6e72e6
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/codeinterpreterapi/session.py b/codeinterpreterapi/session.py
@@ -23,7 +23,7 @@
     RedisChatMessageHistory,
 )
 from langchain.prompts.chat import MessagesPlaceholder
-from langchain.schema import BaseChatMessageHistory, BaseLanguageModel
+from langchain.schema import BaseChatMessageHistory, BaseLanguageModel, SystemMessage
 from langchain.tools import BaseTool, StructuredTool
 
 from codeinterpreterapi.agents import OpenAIFunctionsAgent
@@ -50,13 +50,17 @@ class CodeInterpreterSession:
     def __init__(
         self,
         llm: Optional[BaseLanguageModel] = None,
+        system_message: SystemMessage = code_interpreter_system_message,
+        max_iterations: int = 9,
         additional_tools: list[BaseTool] = [],
         **kwargs,
     ) -> None:
         self.codebox = CodeBox()
         self.verbose = kwargs.get("verbose", settings.VERBOSE)
         self.tools: list[BaseTool] = self._tools(additional_tools)
         self.llm: BaseLanguageModel = llm or self._choose_llm(**kwargs)
+        self.max_iterations = max_iterations
+        self.system_message = system_message
         self.agent_executor: Optional[AgentExecutor] = None
         self.input_files: list[File] = []
         self.output_files: list[File] = []
@@ -148,7 +152,7 @@ def _choose_agent(self) -> BaseSingleActionAgent:
             OpenAIFunctionsAgent.from_llm_and_tools(
                 llm=self.llm,
                 tools=self.tools,
-                system_message=code_interpreter_system_message,
+                system_message=self.system_message,
                 extra_prompt_messages=[
                     MessagesPlaceholder(variable_name="chat_history")
                 ],
@@ -189,7 +193,7 @@ def _history_backend(self) -> BaseChatMessageHistory:
     def _agent_executor(self) -> AgentExecutor:
         return AgentExecutor.from_agent_and_tools(
             agent=self._choose_agent(),
-            max_iterations=9,
+            max_iterations=self.max_iterations,
             tools=self.tools,
             verbose=self.verbose,
             memory=ConversationBufferMemory(