add local storage

.dockerignore

@@ -5,4 +5,5 @@ workspace
 dist
 data
 geckodriver.log
-logs
+logs
+storage

.gitignore

@@ -169,3 +169,4 @@ output.wav
 output
 tmp.png
 
+storage/*

app.py

@@ -6,6 +6,7 @@ import asyncio
 from collections import deque
 import contextlib
 from functools import partial
+import shutil
 import urllib.parse
 from datetime import datetime
 import uuid
@@ -154,58 +155,60 @@ async def create_message(req_model: NewMsg, request: Request):
     """
     Session message stream
     """
-    config = {k.upper(): v for k, v in req_model.config.items()}
-    set_context(config, uuid.uuid4().hex)
-
-    msg_queue = deque()
-    CONFIG.LLM_STREAM_LOG = lambda x: msg_queue.appendleft(x) if x else None
-
-    role = SoftwareCompany()
-    role.recv(message=Message(content=req_model.query))
-    answer = MessageJsonModel(
-        steps=[
-            Sentences(
-                contents=[
-                    Sentence(type=SentenceType.TEXT.value, value=SentenceValue(answer=req_model.query), is_finished=True)
-                ],
-                status=MessageStatus.COMPLETE.value,
-            )
-        ],
-        qa_type=QueryAnswerType.Answer.value,
-    )
-
-    tc_id = 0
-
-    while True:
-        tc_id += 1
-        if request and await request.is_disconnected():
-            return
-        think_result: RoleRun = await role.think()
-        if not think_result:  # End of conversion
-            break
-
-        think_act_prompt = ThinkActPrompt(role=think_result.role.profile)
-        think_act_prompt.update_think(tc_id, think_result)
-        yield think_act_prompt.prompt + "\n\n"
-        task = asyncio.create_task(role.act())
-
-        while not await request.is_disconnected():
-            if msg_queue:
-                think_act_prompt.update_act(msg_queue.pop(), False)
-                yield think_act_prompt.prompt + "\n\n"
-                continue
-
-            if task.done():
+    try:
+        config = {k.upper(): v for k, v in req_model.config.items()}
+        set_context(config, uuid.uuid4().hex)
+
+        msg_queue = deque()
+        CONFIG.LLM_STREAM_LOG = lambda x: msg_queue.appendleft(x) if x else None
+
+        role = SoftwareCompany()
+        role.recv(message=Message(content=req_model.query))
+        answer = MessageJsonModel(
+            steps=[
+                Sentences(
+                    contents=[
+                        Sentence(type=SentenceType.TEXT.value, value=SentenceValue(answer=req_model.query), is_finished=True)
+                    ],
+                    status=MessageStatus.COMPLETE.value,
+                )
+            ],
+            qa_type=QueryAnswerType.Answer.value,
+        )
+
+        tc_id = 0
+
+        while True:
+            tc_id += 1
+            if request and await request.is_disconnected():
+                return
+            think_result: RoleRun = await role.think()
+            if not think_result:  # End of conversion
                 break
 
-            await asyncio.sleep(0.5)
+            think_act_prompt = ThinkActPrompt(role=think_result.role.profile)
+            think_act_prompt.update_think(tc_id, think_result)
+            yield think_act_prompt.prompt + "\n\n"
+            task = asyncio.create_task(role.act())
 
-        act_result = await task
-        think_act_prompt.update_act(act_result)
-        yield think_act_prompt.prompt + "\n\n"
-        answer.add_think_act(think_act_prompt)
-    yield answer.prompt + "\n\n"  # Notify the front-end that the message is complete.
+            while not await request.is_disconnected():
+                if msg_queue:
+                    think_act_prompt.update_act(msg_queue.pop(), False)
+                    yield think_act_prompt.prompt + "\n\n"
+                    continue
 
+                if task.done():
+                    break
+
+                await asyncio.sleep(0.5)
+
+            act_result = await task
+            think_act_prompt.update_act(act_result)
+            yield think_act_prompt.prompt + "\n\n"
+            answer.add_think_act(think_act_prompt)
+        yield answer.prompt + "\n\n"  # Notify the front-end that the message is complete.
+    finally:
+        shutil.rmtree(CONFIG.WORKSPACE_PATH)
 
 default_llm_stream_log = partial(print, end="")
 
@@ -236,10 +239,11 @@ class ChatHandler:
 app = FastAPI()
 
 app.mount(
-    "/static",
-    StaticFiles(directory="./static/", check_dir=True),
+    "/storage",
+    StaticFiles(directory="./storage/"),
     name="static",
 )
+
 app.add_api_route(
     "/api/messages",
     endpoint=ChatHandler.create_message,
@@ -250,13 +254,17 @@ app.add_api_route(
 
 @app.get("/{catch_all:path}")
 async def catch_all(request: Request):
-    if request.url.path == "/":
-        return RedirectResponse(url="/static/index.html")
     if request.url.path.startswith("/api"):
         raise HTTPException(status_code=404)
 
-    new_path = f"/static{request.url.path}"
-    return RedirectResponse(url=new_path)
+    return RedirectResponse(url="/index.html")
+
+
+app.mount(
+    "/",
+    StaticFiles(directory="./static/", html=True),
+    name="static",
+)
 
 
 set_llm_stream_logfunc(llm_stream_log)

config/config.yaml

@@ -120,3 +120,17 @@ RPM: 10
 # PROMPT_FORMAT: json #json or markdown
 
 DISABLE_LLM_PROVIDER_CHECK: true
+STORAGE_TYPE: local # local / s3
+
+# for local storage
+LOCAL_ROOT: "storage"
+LOCAL_BASE_URL: "storage"
+
+
+# for s3 storage
+
+# S3_ACCESS_KEY: ""
+# S3_SECRET_KEY: ""
+# S3_ENDPOINT_URL: ""
+# S3_BUCKET: ""
+# S3_SECURE: false

software_company.py

@@ -62,7 +62,7 @@ class PackProject(Action):
         chunks = []
         async for chunk in AioZipStream(files, chunksize=32768).stream():
             chunks.append(chunk)
-        return await upload_to_s3(b"".join(chunks), key)
+        return await get_download_url(b"".join(chunks), key)
 
 
 class SoftwareCompany(Role):
@@ -345,22 +345,31 @@ class SoftwareCompany(Role):
 async def upload_file_to_s3(filepath: str, key: str):
     async with aiofiles.open(filepath, "rb") as f:
         content = await f.read()
-        return await upload_to_s3(content, key)
-
-
-async def upload_to_s3(content: bytes, key: str):
-    session = get_session()
-    async with session.create_client(
-        "s3",
-        aws_secret_access_key=CONFIG.get("S3_SECRET_KEY"),
-        aws_access_key_id=CONFIG.get("S3_ACCESS_KEY"),
-        endpoint_url=CONFIG.get("S3_ENDPOINT_URL"),
-        use_ssl=CONFIG.get("S3_SECURE"),
-    ) as client:
-        # upload object to amazon s3
-        bucket = CONFIG.get("S3_BUCKET")
-        await client.put_object(Bucket=bucket, Key=key, Body=content)
-        return f"{CONFIG.get('S3_ENDPOINT_URL')}/{bucket}/{key}"
+        return await get_download_url(content, key)
+
+
+async def get_download_url(content: bytes, key: str) -> str:
+    if CONFIG.get("STORAGE_TYPE") == "S3":
+        session = get_session()
+        async with session.create_client(
+            "s3",
+            aws_secret_access_key=CONFIG.get("S3_SECRET_KEY"),
+            aws_access_key_id=CONFIG.get("S3_ACCESS_KEY"),
+            endpoint_url=CONFIG.get("S3_ENDPOINT_URL"),
+            use_ssl=CONFIG.get("S3_SECURE"),
+        ) as client:
+            # upload object to amazon s3
+            bucket = CONFIG.get("S3_BUCKET")
+            await client.put_object(Bucket=bucket, Key=key, Body=content)
+            return f"{CONFIG.get('S3_ENDPOINT_URL')}/{bucket}/{key}"
+    else:
+        storage = CONFIG.get("LOCAL_ROOT", "storage")
+        base_url = CONFIG.get("LOCAL_BASE_URL", "storage")
+        filepath = Path(storage) / key
+        filepath.parent.mkdir(exist_ok=True, parents=True)
+        async with aiofiles.open(filepath, "wb") as f:
+            await f.write(content)
+        return f"{base_url}/{key}"
 
 
 async def main(idea, **kwargs):