refactoring the agent summary generation flow

TheSnowGuru · Jul 2, 2023 · 30c8797 · 30c8797
1 parent 09cf26d
commit 30c8797
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 65 deletions.
diff --git a/superagi/jobs/agent_executor.py b/superagi/jobs/agent_executor.py
@@ -8,6 +8,7 @@
 from superagi.agent.super_agi import SuperAgi
 from superagi.config.config import get_config
 from superagi.helper.encyption_helper import decrypt_data
+from superagi.jobs.resource_summary import ResourceSummarizer
 from superagi.lib.logger import logger
 from superagi.llms.openai import OpenAi
 from superagi.models.agent import Agent
@@ -194,8 +195,8 @@ def execute_next_action(self, agent_execution_id):
             tool = AgentExecutor.create_object(tool, session)
             tools.append(tool)
 
-        resource_summary = self.generate_resource_summary(agent.id, session, model_api_key)
-        resource_summary = resource_summary or parsed_config.get("resource_summary")
+        resource_summary = self.get_agent_resource_summary(agent_id=agent.id, session=session,
+                                                           default_summary=parsed_config.get("resource_summary"))
         tools = self.set_default_params_tools(tools, parsed_config, agent_execution.agent_id,
                                               model_api_key=model_api_key,
                                               resource_description=resource_summary,
@@ -307,40 +308,10 @@ def handle_wait_for_permission(self, agent_execution, spawned_agent, session):
         agent_execution.status = "RUNNING"
         session.commit()
 
-    def generate_resource_summary(self, agent_id: int, session: Session, openai_api_key: str):
-        resources = session.query(Resource).filter(Resource.agent_id == agent_id).order_by(
-            Resource.updated_at.asc()).all()
-        if len(resources) == 0:
-            return
-        # get last resource from agent config
-        agent_last_resource = session.query(AgentConfiguration).filter(AgentConfiguration.agent_id == agent_id,
-                                                                       AgentConfiguration.key == "last_resource").first()
-        if agent_last_resource is not None and \
-                datetime.strptime(agent_last_resource.value, '%Y-%m-%d %H:%M:%S.%f') == resources[-1].updated_at:
-            return
-        texts = [resource.summary for resource in resources if resource.summary is not None]
-        if len(texts) == 0:
-            return
-        if len(texts) > 1:
-            resource_summary = LlamaDocumentSummary().generate_summary_of_texts(texts, openai_api_key)
-        else:
-            resource_summary = texts[0]
-
+    def get_agent_resource_summary(self, agent_id: int, session: Session, default_summary: str):
+        ResourceSummarizer(session=session).generate_agent_summary(agent_id=agent_id)
         agent_config_resource_summary = session.query(AgentConfiguration). \
             filter(AgentConfiguration.agent_id == agent_id,
                    AgentConfiguration.key == "resource_summary").first()
-
-        if agent_config_resource_summary is not None:
-            agent_config_resource_summary.value = resource_summary
-        else:
-            agent_config_resource_summary = AgentConfiguration(agent_id=agent_id, key="resource_summary",
-                                                               value=resource_summary)
-            session.add(agent_config_resource_summary)
-        if agent_last_resource is not None:
-            agent_last_resource.value = str(resources[-1].updated_at)
-        else:
-            agent_last_resource = AgentConfiguration(agent_id=agent_id, key="last_resource",
-                                                     value=str(resources[-1].updated_at))
-            session.add(agent_last_resource)
-        session.commit()
+        resource_summary = agent_config_resource_summary.value if agent_config_resource_summary is not None else default_summary
         return resource_summary
diff --git a/superagi/jobs/resource_summary.py b/superagi/jobs/resource_summary.py
@@ -3,18 +3,20 @@
 from superagi.config.config import get_config
 from superagi.lib.logger import logger
 from superagi.models.agent_config import AgentConfiguration
-from superagi.models.db import connect_db
 from superagi.models.resource import Resource
 from superagi.resource_manager.llama_document_summary import LlamaDocumentSummary
 from superagi.resource_manager.resource_manager import ResourceManager
+from datetime import datetime
 
-engine = connect_db()
-session = sessionmaker(bind=engine)
 
 class ResourceSummarizer:
+    """Class to summarize a resource."""
+
+    def __int__(self, session):
+        self.session = session
 
     @classmethod
-    def add_to_vector_store_and_create_summary(cls, agent_id: int, resource_id: int, documents: list):
+    def add_to_vector_store_and_create_summary(self, agent_id: int, resource_id: int, documents: list):
         """
         Add a file to the vector store and generate a summary for it.
 
@@ -24,7 +26,6 @@ def add_to_vector_store_and_create_summary(cls, agent_id: int, resource_id: int,
             openai_api_key (str): OpenAI API key.
             documents (list): List of documents.
         """
-        db = session()
         try:
             ResourceManager(agent_id).save_document_to_vector_store(documents, str(resource_id))
         except Exception as e:
@@ -34,38 +35,38 @@ def add_to_vector_store_and_create_summary(cls, agent_id: int, resource_id: int,
             summary = LlamaDocumentSummary().generate_summary_of_document(documents)
         except Exception as e:
             logger.error(e)
-        resource = db.query(Resource).filter(Resource.id == resource_id).first()
+        resource = self.session.query(Resource).filter(Resource.id == resource_id).first()
         resource.summary = summary
-        db.commit()
-        resources = db.query(Resource).filter(Resource.agent_id == agent_id).all()
-        summary_texts = [resource.summary for resource in resources if resource.summary is not None]
-        if len(summary_texts) != len(resources):
-            return
-
-        if len(summary_texts) == 1:
-            resource_summary = summary_texts[0]
-        else:
-            resource_summary = LlamaDocumentSummary().generate_summary_of_texts(summary_texts)
+        self.session.commit()
 
-        agent_config_resource_summary = db.query(AgentConfiguration). \
+    def generate_agent_summary(self, agent_id: int) -> str:
+        """Generate a summary of all resources for an agent."""
+        agent_config_resource_summary = self.session.query(AgentConfiguration). \
             filter(AgentConfiguration.agent_id == agent_id,
                    AgentConfiguration.key == "resource_summary").first()
+        resources = self.session.query(Resource).filter(Resource.agent_id == agent_id).all()
+        summary_texts = [resource.summary for resource in resources if resource.summary is not None]
+        if len(summary_texts) != len(resources):
+            return
 
-        agent_last_resource = db.query(AgentConfiguration). \
+        agent_last_resource = self.session.query(AgentConfiguration). \
             filter(AgentConfiguration.agent_id == agent_id,
-                   AgentConfiguration.key == "last_resource").first()
+                   AgentConfiguration.key == "last_resource_id").first()
+        if agent_last_resource is not None and \
+                datetime.strptime(agent_last_resource.value, '%Y-%m-%d %H:%M:%S.%f') == resources[-1].updated_at:
+            return
 
+        resource_summary = LlamaDocumentSummary().generate_summary_of_texts(summary_texts)
         if agent_config_resource_summary is not None:
             agent_config_resource_summary.value = resource_summary
         else:
             agent_config_resource_summary = AgentConfiguration(agent_id=agent_id, key="resource_summary",
                                                                value=resource_summary)
-            db.add(agent_config_resource_summary)
+            self.session.add(agent_config_resource_summary)
         if agent_last_resource is not None:
-            agent_last_resource.value = str(resource.updated_at)
+            agent_last_resource.value = str(self.session.updated_at)
         else:
-            agent_last_resource = AgentConfiguration(agent_id=agent_id, key="last_resource",
-                                                     value=str(resource.updated_at))
-            db.add(agent_last_resource)
-        db.commit()
-        db.close()
+            agent_last_resource = AgentConfiguration(agent_id=agent_id, key="last_resource_id",
+                                                     value=str(resources[-1].updated_at))
+            self.session.add(agent_last_resource)
+        self.session.commit()
diff --git a/superagi/worker.py b/superagi/worker.py
@@ -1,9 +1,13 @@
 from __future__ import absolute_import
+
+from sqlalchemy.orm import sessionmaker
+
 from superagi.lib.logger import logger
 
 from celery import Celery
 
 from superagi.config.config import get_config
+from superagi.models.db import connect_db
 
 redis_url = get_config('REDIS_URL') or 'localhost:6379'
 
@@ -14,7 +18,6 @@
 app.conf.accept_content = ['application/x-python-serialize', 'application/json']
 
 
-
 @app.task(name="execute_agent", autoretry_for=(Exception,), retry_backoff=2, max_retries=5)
 def execute_agent(agent_execution_id: int, time):
     """Execute an agent step in background."""
@@ -28,6 +31,14 @@ def summarize_resource(agent_id: int, resource_id: int,
                        documents: list):
     """Summarize a resource in background."""
     from superagi.jobs.resource_summary import ResourceSummarizer
+    engine = connect_db()
+    Session = sessionmaker(bind=engine)
+    session = Session()
+
     logger.info("Summarize resource:" + str(agent_id) + "," + str(resource_id))
-    ResourceSummarizer.add_to_vector_store_and_create_summary(agent_id=agent_id, resource_id=resource_id,
-                                                              documents=documents)
+    resource_summarizer = ResourceSummarizer(session=session)
+    resource_summarizer.add_to_vector_store_and_create_summary(agent_id=agent_id,
+                                                                                 resource_id=resource_id,
+                                                                                 documents=documents)
+    resource_summarizer.generate_agent_summary(agent_id=agent_id)
+    session.close()
diff --git a/tests/unit_tests/agent_executor/test_generate_resource_summary.py b/tests/unit_tests/agent_executor/test_generate_resource_summary.py
@@ -150,7 +150,7 @@ def test_generate_resource_summary_already_summarized(generate_summary_of_texts)
     session = MockSession()
     openai_api_key = "mock_openai_key"
     agent_id = 1
-    last_resource = MockAgentConfiguration(agent_id=agent_id, key="last_resource", value="2023-07-01 07:47:27.349297")
+    last_resource = MockAgentConfiguration(agent_id=agent_id, key="last_resource_id", value="2023-07-01 07:47:27.349297")
     with patch('superagi.models.agent_config.AgentConfiguration', new=MockAgentConfiguration):
         with patch('superagi.models.resource.Resource', new=MockResource):
             with patch.object(MockSession, 'first', return_value=last_resource):