Merge pull request shroominic#21 from Valian/verbosity-in-session

Verbosity in session
Demolicity · Jul 18, 2023 · 44422c1 · 44422c1
2 parents d871fbf + a31142b
commit 44422c1
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/codeinterpreterapi/session.py b/codeinterpreterapi/session.py
@@ -20,14 +20,15 @@
 
 
 class CodeInterpreterSession:
-    def __init__(self, model=None, openai_api_key=None) -> None:
+    def __init__(self, model=None, openai_api_key=settings.OPENAI_API_KEY, verbose=settings.VERBOSE) -> None:
         self.codebox = CodeBox()
+        self.verbose = verbose
         self.tools: list[StructuredTool] = self._tools()
         self.llm: BaseChatModel = self._llm(model, openai_api_key)
         self.agent_executor: AgentExecutor = self._agent_executor()
         self.input_files: list[File] = []
         self.output_files: list[File] = []
-    
+
     async def astart(self) -> None:
         await self.codebox.astart()
 
@@ -51,10 +52,9 @@ def _llm(self, model: Optional[str] = None, openai_api_key: Optional[str] = None
             model = "gpt-4"
 
         if openai_api_key is None:
-            if settings.OPENAI_API_KEY is None:
-                raise ValueError("OpenAI API key missing.")
-            else:
-                openai_api_key = settings.OPENAI_API_KEY
+            raise ValueError(
+                "OpenAI API key missing. Set OPENAI_API_KEY env variable or pass `openai_api_key` to session."
+            )
 
         return ChatOpenAI(
             temperature=0.03,
@@ -78,13 +78,13 @@ def _agent_executor(self) -> AgentExecutor:
             callbacks=[CodeCallbackHandler(self)],
             max_iterations=9,
             tools=self.tools,
-            verbose=settings.VERBOSE,
+            verbose=self.verbose,
             memory=ConversationBufferMemory(memory_key="memory", return_messages=True),
         )
 
     async def show_code(self, code: str) -> None:
         """Callback function to show code to the user."""
-        if settings.VERBOSE:
+        if self.verbose:
             print(code)
 
     def run_handler(self, code: str):
@@ -113,7 +113,7 @@ async def arun_handler(self, code: str):
                     return f"{package.group(1)} was missing but got installed now. Please try again."
             else: pass
                 # TODO: preanalyze error to optimize next code generation
-            if settings.VERBOSE:
+            if self.verbose:
                 print("Error:", output.content)
 
         elif modifications := await get_file_modifications(code, self.llm):
@@ -170,7 +170,7 @@ async def generate_response(
             response = await self.agent_executor.arun(input=user_request.content)
             return await self.output_handler(response)
         except Exception as e:
-            if settings.VERBOSE:
+            if self.verbose:
                 import traceback
 
                 traceback.print_exc()
@@ -186,10 +186,10 @@ async def generate_response(
 
     async def is_running(self) -> bool:
         return await self.codebox.astatus() == "running"
-    
+
     async def astop(self) -> None:
         await self.codebox.astop()
-    
+
     async def __aenter__(self) -> "CodeInterpreterSession":
         await self.astart()
         return self