Autogen skill and server support (#87)

* add autogen manager * add autogen skill * autogen server * cleanup --------- Co-authored-by: daviddai <wdai98@163.com>
AI-Citizen · Oct 30, 2023 · 4b22b3e · 4b22b3e
1 parent cbe00be
commit 4b22b3e
Show file tree

Hide file tree

Showing 15 changed files with 598 additions and 4 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -34,3 +34,4 @@ rq==1.15.1
 eventlet==0.33.3
 tiktoken==0.5.1
 text-generation==0.6.1
+pyautogen==0.1.13
diff --git a/solidgpt/src/api/api.py b/solidgpt/src/api/api.py
@@ -54,6 +54,7 @@ def set_name(self, new_name):
 Initializer()
 # orchestration = Orchestration()
 uploaded_repo_map: dict = {}
+autogen_task_map: dict = {}
 graph_result_map: dict[str, GraphResult] = {}
 graph_stage_map: dict = {}
 
@@ -168,9 +169,10 @@ async def generate_tech_solution(body: dict = Body(...)):
         "current_work_name": "tech solution"
     }, status_code=200)
 
+
 @app.post("/repochat")
 async def repo_chat(body: dict = Body(...)):
-    # Enqueue the background task: tech solution
+    # Enqueue the background task: repo chat
     logging.info("celery task: repo chat graph")
     graph_id = str(uuid.uuid4())
     onboarding_id = body['onboarding_id']
@@ -189,6 +191,51 @@ async def repo_chat(body: dict = Body(...)):
     }, status_code=200)
 
 
+@app.post("/autogenanalysis")
+async def autogen_analysis(body: dict = Body(...)):
+    # Enqueue the background task: autogen analysis
+    logging.info("celery task: autogen analysis graph")
+
+    onboarding_id = body['onboarding_id']
+    openai_key = body['openai_key']
+    requirement = body['requirement']
+    task_id = body['task_id']
+    is_new_session = int(body['is_new_session'])
+
+    if is_new_session:
+        graph_id = str(uuid.uuid4())
+        result = celery_task_autogen_analysis_graph.apply_async(args=[
+            openai_key, requirement, onboarding_id, graph_id])
+        task_id = result.id
+        autogen_task_map[task_id] = result
+
+        return JSONResponse(content={
+            "message": f"New autogen analysis graph...",
+            "task_id": task_id,
+            "is_final": True,
+            "status": 1,
+            "current_work_name": "autogen analysis"
+        }, status_code=200)
+    else:
+        task = autogen_task_map.get(task_id)
+        if task is None:
+            return JSONResponse(content={
+                "message": f"Invalid Autogen Analysis graph.",
+                "task_id": task_id,
+                "is_final": True,
+                "status": 2,
+                "current_work_name": "autogen analysis"
+            }, status_code=200)
+        r.lpush(task_id, requirement)
+        return JSONResponse(content={
+            "message": f"Continuing autogen analysis graph...",
+            "task_id": task_id,
+            "is_final": True,
+            "status": 3,
+            "current_work_name": "autogen analysis"
+        }, status_code=200)
+
+
 @app.post("/uploadrepo")
 async def upload_repo(body: dict = Body(...)):
     # Store the file to local storage
@@ -204,6 +251,33 @@ async def upload_repo(body: dict = Body(...)):
     }, status_code=200)
 
 
+@app.post("/status/autogen")
+async def get_autogen_status(body: dict = Body(...)):
+    task_id: str = body['task_id']
+    celery_task_result = autogen_task_map.get(task_id)
+
+    if celery_task_result is None:
+        return JSONResponse(content={
+            "message": "status: Error, not exist or not started",
+            "task_id": task_id,
+            "status": 1,
+            "result": ""
+        }, status_code=200)
+    if celery_task_result.ready():
+        return JSONResponse(content={
+            "message": "status: Error, autogen task cannot be finished",
+            "task_id": task_id,
+            "status": 2,
+            "result": ""
+        }, status_code=200)
+    return JSONResponse(content={
+        "message": "status: autogen task result",
+        "task_id": task_id,
+        "status": 3,
+        "result": celery_task_result.info
+    }, status_code=200)
+
+
 @app.post("/status/upload")
 async def get_upload_status(body: dict = Body(...)):
     repo_upload_id: str = body['upload_id']

diff --git a/solidgpt/src/api/celery_tasks.py b/solidgpt/src/api/celery_tasks.py
@@ -1,5 +1,6 @@
 import base64
 import time
+import uuid
 
 from celery import Celery
 import shutil
@@ -13,12 +14,14 @@
 from solidgpt.src.workskill.skills.load_repo import LoadRepo
 from solidgpt.src.workgraph.graph import *
 from solidgpt.src.manager.initializer import *
+import redis
 
 app = Celery('celery_tasks',
              BROKER_URL='redis://localhost:6379/0',  # Using Redis as the broker
              CELERY_RESULT_BACKEND='redis://localhost:6379/1'
              )
 app.autodiscover_tasks(['solidgpt.src.api'])
+r = redis.Redis()
 Initializer()
 
 
@@ -104,12 +107,48 @@ def celery_task_tech_solution_graph(openai_key, requirement, onboarding_id, grap
     text_result = g.display_result.get_result()
     return text_result
 
+
 @app.task
 def celery_task_repo_chat_graph(openai_key, requirement, onboarding_id, graph_id):
-    logging.info("celery task: tech solution graph")
+    logging.info("celery task: repo chat graph")
     openai.api_key = openai_key
     g = build_repo_chat_graph(requirement, onboarding_id, graph_id)
     g.init_node_dependencies()
     g.execute()
     text_result = g.display_result.get_result()
     return text_result
+
+
+@app.task(bind=True)
+def celery_task_autogen_analysis_graph(self, openai_key, requirement, onboarding_id, graph_id):
+    logging.info("celery task: autogen analysis graph")
+    self.update_state(
+        state='PROGRESS',
+        meta={'result': "", 'state_id': ""}
+    )
+
+    def autogen_message_input_callback():
+        data = r.blpop(self.request.id)
+        if data:
+            # Extracting UUID and poem text from the tuple
+            task_id, poem_bytes = data
+
+            # Converting bytes to string
+            poem_text = poem_bytes.decode()
+            print(poem_text)  # for server debug
+            return poem_text
+        return ""
+
+    def autogen_update_result_callback(result):
+        self.update_state(
+            state='PROGRESS',
+            meta={'result': result, 'state_id': str(uuid.uuid4())}
+        )
+
+    openai.api_key = openai_key
+    g = build_autogen_analysis_graph(requirement, onboarding_id, graph_id,
+                                     autogen_message_input_callback, autogen_update_result_callback)
+    g.init_node_dependencies()
+    g.execute()
+
+    return ""
diff --git a/solidgpt/src/constants.py b/solidgpt/src/constants.py
@@ -11,4 +11,5 @@
 SKILL_NAME_SUMMARY_FILE = "Summary file"
 SKILL_NAME_ANALYSIS_PRODUCT = "Analysis Product"
 SKILL_NAME_PROVIDE_TECH_SOLUTION = "Provide Tech Solution"
-SKILL_NAME_REPO_CHAT = "Repo Chat"
+SKILL_NAME_REPO_CHAT = "Repo Chat"
+SKILL_NAME_AUTOGEN_ANALYSIS = "AutoGen analysis"