From 5c7eeecacb38fdffd1207db4cf62861fcc26500d Mon Sep 17 00:00:00 2001
From: jinno <klvx01@gmail.com>
Date: Mon, 2 Sep 2024 22:28:16 +0900
Subject: [PATCH] fix: update brain.py and CodeInterpreterIntermediateResult

---
 src/codeinterpreterapi/brain/brain.py |  2 ++
 src/codeinterpreterapi/schema.py      | 11 +++++++++++
 2 files changed, 13 insertions(+)

diff --git a/src/codeinterpreterapi/brain/brain.py b/src/codeinterpreterapi/brain/brain.py
index 5e5496ad..1ae8eb33 100644
--- a/src/codeinterpreterapi/brain/brain.py
+++ b/src/codeinterpreterapi/brain/brain.py
@@ -149,6 +149,8 @@ def run(
         self.update_agent_score()
         if isinstance(output, str):
             output = CodeInterpreterIntermediateResult(context=output)
+        elif not isinstance(output, CodeInterpreterIntermediateResult):
+            output = CodeInterpreterIntermediateResult(context=str(output))
         return output
 
     def _set_output_llm_result(
diff --git a/src/codeinterpreterapi/schema.py b/src/codeinterpreterapi/schema.py
index 3171ab80..556e8c58 100644
--- a/src/codeinterpreterapi/schema.py
+++ b/src/codeinterpreterapi/schema.py
@@ -135,6 +135,17 @@ class CodeInterpreterIntermediateResult(BaseModel):
     iteration_count: int = Field(default=0, description="現在の反復回数")
     max_iterations: int = Field(default=10, description="最大反復回数")
 
+    def __str__(self) -> str:
+        return self.__repr__()
+
+    def __repr__(self) -> str:
+        return f"""CodeInterpreterIntermediateResult(
+            thoughts={self.thoughts},
+            context={self.context},
+            code={self.code},
+            log={self.log}
+            confidence={self.confidence})"""
+
 
 class CodeInterpreterResponse(AIMessage):
     """