go

edenartlab · Jan 9, 2025 · ac52749 · ac52749
1 parent 6f7478e
commit ac52749
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/eve/llm.py b/eve/llm.py
@@ -17,6 +17,7 @@
 
 from . import sentry_sdk
 from .tool import Tool
+from .task import Creation
 from .user import User
 from .agent import Agent
 from .thread import UserMessage, AssistantMessage, ToolCall, Thread
@@ -574,15 +575,34 @@ async def async_prompt_thread(
                 result = await tool.async_wait(task)
                 thread.update_tool_call(assistant_message.id, t, result)
 
-                # yield update
+                # task completed
                 if result["status"] == "completed":
+
+                    # make a Creation
+                    name = task.args.get("prompt") or task.args.get("text_input")
+                    filename = result.get("output", [{}])[0].get("filename")
+                    media_attributes = result.get("output", [{}])[0].get("mediaAttributes")
+                    if filename and media_attributes:
+                        new_creation = Creation(
+                            user=task.user,
+                            requester=task.requester,
+                            task=task.id,
+                            tool=task.tool,
+                            filename=filename,
+                            mediaAttributes=media_attributes,
+                            name=name
+                        )
+                        new_creation.save()
+
+                    # yield update
                     yield ThreadUpdate(
                         type=UpdateType.TOOL_COMPLETE,
                         tool_name=tool_call.tool,
                         tool_index=t,
                         result=result,
                     )
                 else:
+                    # yield error
                     yield ThreadUpdate(
                         type=UpdateType.ERROR,
                         tool_name=tool_call.tool,

diff --git a/eve/task.py b/eve/task.py
@@ -140,7 +140,8 @@ async def _task_handler(func, *args, **kwargs):
                 result = eden_utils.upload_result(result, save_thumbnails=True, save_blurhash=True)
 
                 for output in result["output"]:
-                    name = preprocess_result.get("name") or task_args.get("prompt") or args.get("text_input")
+                    # name = preprocess_result.get("name") or task_args.get("prompt") or args.get("text_input")
+                    name = task_args.get("prompt") or args.get("text_input")
                     if not name:
                         name = args.get("interpolation_prompts") or args.get("interpolation_texts")
                         if name: