Merge pull request #77 from shroominic/chat_history_backends

Chat history backends + Session ID management

Author: shroominic

.env.example

@@ -1,11 +1,20 @@
+# (set True to enable logging)
+VERBOSE=False
+
 # (required)
 OPENAI_API_KEY=
+# ANTHROPIC_API_KEY=
+
 # (optional, required for production)
 # CODEBOX_API_KEY=
-# (set True to enable logging)
-VERBOSE=False
+
 # (optional, required for Azure OpenAI)
 # OPENAI_API_TYPE=azure
 # OPENAI_API_VERSION=2023-07-01-preview
 # OPENAI_API_BASE=
 # DEPLOYMENT_NAME=
+
+# (optional, [codebox, postgres or redis])
+# HISTORY_BACKEND=postgres
+# REDIS_URL=
+# POSTGRES_URL=

codeinterpreterapi/chat_history.py

@@ -0,0 +1,62 @@
+import asyncio
+import json
+from typing import List
+
+from codeboxapi import CodeBox  # type: ignore
+from langchain.schema import BaseChatMessageHistory
+from langchain.schema.messages import BaseMessage, messages_from_dict, messages_to_dict
+
+
+# TODO: This is probably not efficient, but it works for now.
+class CodeBoxChatMessageHistory(BaseChatMessageHistory):
+    """
+    Chat message history that stores history inside the codebox.
+    """
+
+    def __init__(self, codebox: CodeBox):
+        self.codebox = codebox
+
+        if "history.json" not in [f.name for f in self.codebox.list_files()]:
+            name, content = "history.json", b"{}"
+            if (loop := asyncio.get_event_loop()).is_running():
+                loop.create_task(self.codebox.aupload(name, content))
+            else:
+                self.codebox.upload(name, content)
+
+    @property
+    def messages(self) -> List[BaseMessage]:  # type: ignore
+        """Retrieve the messages from the codebox"""
+        msgs = (
+            messages_from_dict(json.loads(file_content.decode("utf-8")))
+            if (
+                file_content := (
+                    loop.run_until_complete(self.codebox.adownload("history.json"))
+                    if (loop := asyncio.get_event_loop()).is_running()
+                    else self.codebox.download("history.json")
+                ).content
+            )
+            else []
+        )
+        return msgs
+
+    def add_message(self, message: BaseMessage) -> None:
+        """Append the message to the record in the local file"""
+        print("Current messages: ", self.messages)
+        messages = messages_to_dict(self.messages)
+        print("Adding message: ", message)
+        messages.append(messages_to_dict([message])[0])
+        name, content = "history.json", json.dumps(messages).encode("utf-8")
+        if (loop := asyncio.get_event_loop()).is_running():
+            loop.create_task(self.codebox.aupload(name, content))
+        else:
+            self.codebox.upload(name, content)
+        print("New messages: ", self.messages)
+
+    def clear(self) -> None:
+        """Clear session memory from the local file"""
+        print("Clearing history CLEARING HISTORY")
+        code = "import os; os.remove('history.json')"
+        if (loop := asyncio.get_event_loop()).is_running():
+            loop.create_task(self.codebox.arun(code))
+        else:
+            self.codebox.run(code)

codeinterpreterapi/config.py

@@ -14,8 +14,12 @@ class CodeInterpreterAPISettings(BaseSettings):
 
     VERBOSE: bool = False
 
-    CODEBOX_API_KEY: Optional[str] = None
     OPENAI_API_KEY: Optional[str] = None
+    CODEBOX_API_KEY: Optional[str] = None
+
+    HISTORY_BACKEND: Optional[str] = None
+    REDIS_URL: str = "redis://localhost:6379"
+    POSTGRES_URL: str = "postgresql://postgres:postgres@localhost:5432/postgres"
 
 
 settings = CodeInterpreterAPISettings()

codeinterpreterapi/schema/file.py

@@ -9,6 +9,8 @@ class File(BaseModel):
 
     @classmethod
     def from_path(cls, path: str):
+        if not path.startswith("/"):
+            path = f"./{path}"
         with open(path, "rb") as f:
             path = path.split("/")[-1]
             return cls(name=path, content=f.read())
@@ -33,6 +35,8 @@ async def afrom_url(cls, url: str):
                 return cls(name=url.split("/")[-1], content=await r.read())
 
     def save(self, path: str):
+        if not path.startswith("/"):
+            path = f"./{path}"
         with open(path, "wb") as f:
             f.write(self.content)
 

codeinterpreterapi/session.py

@@ -1,10 +1,10 @@
 import base64
 import re
 import traceback
-import uuid
 from io import BytesIO
 from os import getenv
 from typing import Optional
+from uuid import UUID, uuid4
 
 from codeboxapi import CodeBox  # type: ignore
 from codeboxapi.schema import CodeBoxOutput  # type: ignore
@@ -17,8 +17,13 @@
 from langchain.chat_models import AzureChatOpenAI, ChatAnthropic, ChatOpenAI
 from langchain.chat_models.base import BaseChatModel
 from langchain.memory import ConversationBufferMemory
+from langchain.memory.chat_message_histories import (
+    ChatMessageHistory,
+    PostgresChatMessageHistory,
+    RedisChatMessageHistory,
+)
 from langchain.prompts.chat import MessagesPlaceholder
-from langchain.schema.language_model import BaseLanguageModel
+from langchain.schema import BaseChatMessageHistory, BaseLanguageModel
 from langchain.tools import BaseTool, StructuredTool
 
 from codeinterpreterapi.agents import OpenAIFunctionsAgent
@@ -28,6 +33,7 @@
     get_file_modifications,
     remove_download_link,
 )
+from codeinterpreterapi.chat_history import CodeBoxChatMessageHistory
 from codeinterpreterapi.config import settings
 from codeinterpreterapi.parser import CodeAgentOutputParser, CodeChatAgentOutputParser
 from codeinterpreterapi.prompts import code_interpreter_system_message
@@ -47,20 +53,35 @@ def __init__(
         additional_tools: list[BaseTool] = [],
         **kwargs,
     ) -> None:
-        self.codebox = CodeBox(**kwargs)
+        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.agent_executor: AgentExecutor = self._agent_executor()
+        self.agent_executor: Optional[AgentExecutor] = None
         self.input_files: list[File] = []
         self.output_files: list[File] = []
         self.code_log: list[tuple[str, str]] = []
 
+    @classmethod
+    def from_id(cls, session_id: UUID, **kwargs) -> "CodeInterpreterSession":
+        session = cls(**kwargs)
+        session.codebox = CodeBox.from_id(session_id)
+        session.agent_executor = session._agent_executor()
+        return session
+
+    @property
+    def session_id(self) -> Optional[UUID]:
+        return self.codebox.session_id
+
     def start(self) -> SessionStatus:
-        return SessionStatus.from_codebox_status(self.codebox.start())
+        status = SessionStatus.from_codebox_status(self.codebox.start())
+        self.agent_executor = self._agent_executor()
+        return status
 
     async def astart(self) -> SessionStatus:
-        return SessionStatus.from_codebox_status(await self.codebox.astart())
+        status = SessionStatus.from_codebox_status(await self.codebox.astart())
+        self.agent_executor = self._agent_executor()
+        return status
 
     def _tools(self, additional_tools: list[BaseTool]) -> list[BaseTool]:
         return additional_tools + [
@@ -108,15 +129,15 @@ def _choose_llm(
                     openai_api_key=openai_api_key,
                     max_retries=3,
                     request_timeout=60 * 3,
-                )
+                )  # type: ignore
             else:
                 return ChatOpenAI(
                     temperature=0.03,
                     model=model,
                     openai_api_key=openai_api_key,
                     max_retries=3,
                     request_timeout=60 * 3,
-                )
+                )  # type: ignore
         elif "claude" in model:
             return ChatAnthropic(model=model)
         else:
@@ -148,14 +169,33 @@ def _choose_agent(self) -> BaseSingleActionAgent:
             )
         )
 
+    def _history_backend(self) -> BaseChatMessageHistory:
+        return (
+            CodeBoxChatMessageHistory(codebox=self.codebox)
+            if settings.HISTORY_BACKEND == "codebox"
+            else RedisChatMessageHistory(
+                session_id=str(self.session_id),
+                url=settings.REDIS_URL,
+            )
+            if settings.HISTORY_BACKEND == "redis"
+            else PostgresChatMessageHistory(
+                session_id=str(self.session_id),
+                connection_string=settings.POSTGRES_URL,
+            )
+            if settings.HISTORY_BACKEND == "postgres"
+            else ChatMessageHistory()
+        )
+
     def _agent_executor(self) -> AgentExecutor:
         return AgentExecutor.from_agent_and_tools(
             agent=self._choose_agent(),
             max_iterations=9,
             tools=self.tools,
             verbose=self.verbose,
             memory=ConversationBufferMemory(
-                memory_key="chat_history", return_messages=True
+                memory_key="chat_history",
+                return_messages=True,
+                chat_memory=self._history_backend(),
             ),
         )
 
@@ -178,7 +218,7 @@ def _run_handler(self, code: str):
             raise TypeError("Expected output.content to be a string.")
 
         if output.type == "image/png":
-            filename = f"image-{uuid.uuid4()}.png"
+            filename = f"image-{uuid4()}.png"
             file_buffer = BytesIO(base64.b64decode(output.content))
             file_buffer.name = filename
             self.output_files.append(File(name=filename, content=file_buffer.read()))
@@ -225,7 +265,7 @@ async def _arun_handler(self, code: str):
             raise TypeError("Expected output.content to be a string.")
 
         if output.type == "image/png":
-            filename = f"image-{uuid.uuid4()}.png"
+            filename = f"image-{uuid4()}.png"
             file_buffer = BytesIO(base64.b64decode(output.content))
             file_buffer.name = filename
             self.output_files.append(File(name=filename, content=file_buffer.read()))
@@ -349,6 +389,7 @@ def generate_response_sync(
         user_request = UserRequest(content=user_msg, files=files)
         try:
             self._input_handler(user_request)
+            assert self.agent_executor, "Session not initialized."
             response = self.agent_executor.run(input=user_request.content)
             return self._output_handler(response)
         except Exception as e:
@@ -392,6 +433,7 @@ async def agenerate_response(
         user_request = UserRequest(content=user_msg, files=files)
         try:
             await self._ainput_handler(user_request)
+            assert self.agent_executor, "Session not initialized."
             response = await self.agent_executor.arun(input=user_request.content)
             return await self._aoutput_handler(response)
         except Exception as e:

examples/chat_history_backend.py

@@ -0,0 +1,33 @@
+import os
+
+os.environ["HISTORY_BACKEND"] = "redis"
+os.environ["REDIS_HOST"] = "redis://localhost:6379"
+
+from codeinterpreterapi import CodeInterpreterSession  # noqa: E402
+
+
+def main():
+    session_id = None
+
+    session = CodeInterpreterSession()
+    session.start()
+
+    print("Session ID:", session.session_id)
+    session_id = session.session_id
+
+    response = session.generate_response_sync("Plot the bitcoin chart of 2023 YTD")
+    response.show()
+
+    del session
+
+    assert session_id is not None
+    session = CodeInterpreterSession.from_id(session_id)
+
+    response = session.generate_response_sync("Now for the last 5 years")
+    response.show()
+
+    session.stop()
+
+
+if __name__ == "__main__":
+    main()