Add memory

schema_agents/teams/image_analysis_hub/data_engineer.py

@@ -4,7 +4,8 @@
 from typing import List, Union
 
 from schema_agents.role import Role
-from schema_agents.schema import Message
+from schema_agents.schema import Message, MemoryChunk
+from schema_agents.memory.long_term_memory import LongTermMemory
 from schema_agents.tools.code_interpreter import create_mock_client
 
 from schema_agents.teams.image_analysis_hub.schemas import (PythonFunctionScript, PythonFunctionScriptChanges, Change,
@@ -83,7 +84,8 @@ async def fix_code(
     )
     response = await role.aask(
         python_function_with_line_number,
-        Union[PythonFunctionScriptChanges, PythonFunctionScript],
+        # Union[PythonFunctionScriptChanges, PythonFunctionScript],
+        PythonFunctionScriptChanges,
         prompt=prompt,
     )
     if isinstance(response, PythonFunctionScript):
@@ -117,6 +119,12 @@ async def fix_code(
         )
         return await fix_code(role, client, python_function, output_summary)
 
+    # add the memory of experience
+    experience = response.experience
+    if experience:
+        error_memo = MemoryChunk(index=experience.summary, content=experience, category='error')
+        role.long_term_memory.add(error_memo)
+
     return PythonFunctionScript(
         function_names=python_function.function_names,
         docstring=python_function.docstring,
@@ -129,6 +137,13 @@ async def fix_code(
 async def generate_code(
     req: SoftwareRequirement, role: Role
 ) -> PythonFunctionScript:
+    # retrieve req-related memories from long term memory
+    memories = role.long_term_memory.retrieve(req.original_requirements, filter={"category": "error"})
+    if memories:
+        # concatenate the error experiences
+        error = "\n".join([m.content.error for m in memories])
+        req.additional_notes += f"\nPlease avoid the following error: {error}" # TODO, what to retreive from the error memory?
+
     return await role.aask(req, PythonFunctionScript)
 
 INSTALL_SCRIPT = """
@@ -174,11 +189,20 @@ async def test_run_python_function(
         )
     return python_function
 
+def create_long_term_memory():
+    memory = LongTermMemory()
+    role_id = 'data_engineer'
+    memory.recover_memory(role_id)
+    memory.clean() # this is for test purpose, to avoid the memory to be accumulated
+
+    return memory
+
 def create_data_engineer(client=None):
     async def develop_python_functions(
         req: SoftwareRequirement, role: Role
     ) -> PythonFunctionScript:
         """Develop python functions based on software requirements."""
+
         # if isinstance(req, SoftwareRequirement):
         func = await generate_code(req, role)
         try:
@@ -187,6 +211,7 @@ async def develop_python_functions(
             req.additional_notes += f"\nPlease avoid the following error: {exp}"
             func = await generate_code(req, role)
             func = await test_run_python_function(role, client, req.id, func)
+
         return func
         # else:
         #     if client:
@@ -204,6 +229,7 @@ async def develop_python_functions(
         goal="Develop the python function script according to the software requirement, ensuring that it fulfills the desired functionality. Implement necessary algorithms, handle data processing, and write tests to validate the correctness of the function.",
         constraints=None,
         actions=[develop_python_functions],
+        long_term_memory=create_long_term_memory(),
     )
     return data_engineer
 

schema_agents/teams/image_analysis_hub/microscopist.py

@@ -3,23 +3,23 @@
 
 from pydantic import BaseModel, Field
 from schema_agents.role import Role
-from schema_agents.schema import Message
+from schema_agents.schema import Message, MemoryChunk
+from schema_agents.memory.long_term_memory import LongTermMemory
 from schema_agents.tools.code_interpreter import create_mock_client
 from schema_agents.utils.common import EventBus
 from schema_agents.logs import logger
+from schema_agents.teams.image_analysis_hub.schemas import ExperienceMemory, FunctionMemory
+
 class MicroscopeControlRequirements(BaseModel):
     """Requirements for controlling the microscope and acquire images."""
     path: str = Field(default="", description="save images path")
     timeout: float = Field(default=0.0, description="timeout")
     query: str = Field(default="", description="user's original request")
     plan: str = Field(default="", description="plan for control microscope and acquiring images")
+
 
 class MultiDimensionalAcquisitionScript(BaseModel):
-    """Python script for simple and complex multi-dimensional acquisition.
-    In the script, you can use the following functions to control the microscope:
-    - `microscope_move({'x': 0.0, 'y': 0.0, 'z': 0.0})` # x, y, z are in microns
-    - `microscope_snap({'path': './images', 'exposure': 0.0})` # path is the path to save the image, exposure is in seconds
-    """
+    """Python script for simple and complex multi-dimensional acquisition. Use the provided 'Experiences' to generate the script."""
     script: str = Field(default="", description="Script for acquiring multi-dimensional images")
     explanation: str = Field(default="", description="Brief explanation for the script")
     timeout: float = Field(default=0.0, description="a reasonable timeout for executing the script")
@@ -30,32 +30,72 @@ class ExecutionResult(BaseModel):
     outputs: List[Dict[str, Any]] = Field(default=[], description="Outputs of executing the script")
     traceback: Optional[str] = Field(default=None, description="Traceback of executing the script")
 
-INIT_SCRIPT = """
-def microscope_move(position):
-    print(f"===> Moving to: {position}")
 
-def microscope_snap(config):
-    print(f"===> Snapped an image with exposure {config['exposure']} and saved to: { config['path']}")
-"""
+class FunctionDescription(BaseModel):
+    """Description of a function."""
+    name: str = Field(description="Name of the function")
+    docstring: str = Field(description="Docstring of the function")
+    args: List[str] = Field(default=[], description="Arguments of the function with type annotation")
+
+
+class ScriptGenerationContext(BaseModel):
+    """The context for generating the script."""
+    registed_functions: List[FunctionDescription] = Field(default=[], description="A list of registed functions which can be used directly in the script.")
+    experience_memory: List[ExperienceMemory] = Field(default=[], description="A list of experience memories to help create the script and avoid common mistakes.")
+
+def create_long_term_memory():
+    memory = LongTermMemory()
+    role_id = 'bio'
+    memory.recover_memory(role_id)
+    memory.clean()
+
+    function_move = FunctionMemory(function_name='microscope_move', code="""def microscope_move(position):
+        print(f"===> Moving to: {position}")""", lang='python', args=['position:Tuple[float,float,float]'], docstring='Move the microscope to the given position.')
+    function_snap = FunctionMemory(function_name='microscope_snap', code="""def microscope_snap(config):
+        print(f"===> Snapped an image with exposure {config['exposure']} and saved to: { config['path']}")""", lang='python', args=['config:Dict[str,Any]'], 
+        docstring='Snap an image with the given configuration. The input config should contain the exposure time (key=exposure) and the path (key=path) to save the image.')
+
+    exp = ExperienceMemory(summary='Microscope move restriction', keypoints='Make sure each movement on microscope is larger than 5nm.')
+
+    exp_memo = MemoryChunk(index='Experience for running microscope_move function', content=exp, category='experience')
+    memory.add(exp_memo)
+    fun1_memo = MemoryChunk(index='microscope move python function',content=function_move, category='function')
+    memory.add(fun1_memo)
+    fun2_memo = MemoryChunk(index='microscope snap python function',content=function_snap, category='function')
+    memory.add(fun2_memo)
+    memories = memory.recover_memory(role_id)
+    return memory
 
 
 class Microscope():
     def __init__(self, client):
         self.client = client
-        self.initialized = False
+        self.initialized = True
 
     async def plan(self, query: str=None, role: Role=None) -> MicroscopeControlRequirements:
         """Make a plan for image acquisition tasks."""
         return await role.aask(query, MicroscopeControlRequirements)
 
     async def multi_dimensional_acquisition(self, config: MicroscopeControlRequirements=None, role: Role=None) -> ExecutionResult:
         """Perform image acquisition by using Python script."""
-        if not self.initialized:
-            await self.client.executeScript({"script": INIT_SCRIPT})
-            self.initialized = True
+
+        function_memories = role.long_term_memory.retrieve("microscope related functions", filter={"category": "function"})
+        function_list = []
+        for memory in function_memories:
+            script = memory.content.code
+            function_list.append(FunctionDescription(name=memory.content.function_name, docstring=memory.content.docstring, args=memory.content.args))
+            await self.client.executeScript({"script": script})
+        experiences = role.long_term_memory.retrieve("microscope related function", filter={"category": "experience"})
+        exp_list = []
+        for exp in experiences:
+            exp_list.append(exp.content)
+        context = ScriptGenerationContext(registed_functions=function_list, experience_memory=exp_list)
         print("Acquiring images in multiple dimensions: " + str(config))
-        controlScript = await role.aask(config, MultiDimensionalAcquisitionScript)
+
+        inputs = [context, config, """Make sure each movement on the microscope stage is larger than 5nm."""]
+        controlScript = await role.aask(inputs, MultiDimensionalAcquisitionScript)
         result = await self.client.executeScript({"script": controlScript.script, "timeout": controlScript.timeout})
+
         return ExecutionResult(
             status=result['status'],
             outputs=result['outputs'],
@@ -65,16 +105,18 @@ async def multi_dimensional_acquisition(self, config: MicroscopeControlRequireme
 def create_microscopist(client=None):
     if not client:
         client = create_mock_client()
+
     microscope = Microscope(client)
     microscopist = Role(
         name="Thomas",
         profile="Microscopist",
         goal="Acquire images from the microscope based on user's requests.",
         constraints=None,
-        actions=[microscope.multi_dimensional_acquisition],
+        actions=[microscope.plan, microscope.multi_dimensional_acquisition],
     )
     return microscopist
 
+
 async def main():
     client = create_mock_client()
     event_bus = EventBus("microscopist")
@@ -87,9 +129,11 @@ async def main():
         goal="Acquire images from the microscope based on user's requests.",
         constraints=None,
         actions=[microscope.plan, microscope.multi_dimensional_acquisition],
+        long_term_memory=create_long_term_memory(),
         event_bus=event_bus
     )
 
+
     messages = await ms.handle(Message(content="acquire an image and save to /tmp/img.png", role="User"))
     assert len(messages) == 2
     assert isinstance(messages[-1].data, ExecutionResult)
@@ -108,4 +152,5 @@ async def main():
 
 if __name__ == "__main__":
     asyncio.run(main())
+    # create_memory_storage()
 

schema_agents/teams/image_analysis_hub/schemas.py

@@ -45,13 +45,8 @@ class FunctionMemory(BaseModel):
     function_name: str = Field(default="", description="Function name")
     code: str = Field(default="", description="original code of the function")
     lang: str = Field(default="", description="function language")
-    args: List[str] = Field(default=[], description="arguments of the function")
-
-class ErrorMemory(BaseModel):
-    """Experience of making errors to be saved in the long term memory."""    
-    error: str = Field(default="", description="Error description")
-    cause_by: str = Field(default="", description="Cause of the error")
-    solution: str = Field(default="", description="Solution to fix the error")
+    args: List[str] = Field(default=[], description="arguments of the function, with type annotation")
+    docstring: Optional[str] = Field(default=None, description="docstring of the function")
 
 class ExperienceMemory(BaseModel):
     """Experience to be saved in the long term memory."""