superagent-ai · homanp · Feb 25, 2024 · Feb 21, 2024 · Feb 21, 2024 · Feb 21, 2024
diff --git a/libs/superagent/app/agents/llm.py b/libs/superagent/app/agents/llm.py
@@ -51,7 +51,7 @@ async def _set_tools_return_direct(self):
 
     async def init(self):
         self.agent_config.type = AgentType.SUPERAGENT
-        self.agent_config.llmModel = LLM_REVERSE_MAPPING.get("gpt-4-0613")
+        self.agent_config.llmModel = LLM_REVERSE_MAPPING.get("gpt-3.5-turbo-0125")
         self.agent_id = self.agent_config.id
 
         await self._set_llm()
@@ -85,7 +85,7 @@ async def ainvoke(self, input, *_, **kwargs):
                         input=input
                     )
 
-                model = agent_config.metadata.get("model", "gpt-3.5-turbo")
+                model = agent_config.metadata.get("model", "gpt-3.5-turbo-0125")
                 prompt = agent_config.prompt
                 api_key = agent_config.llms[0].llm.apiKey
 

diff --git a/libs/superagent/app/api/workflows.py b/libs/superagent/app/api/workflows.py
@@ -263,8 +263,11 @@ async def send_message() -> AsyncIterable[str]:
                         yield f"id: {workflow_step['agent_name']}\ndata: {token}\n\n"
 
                 await task
-                workflow_result = task.result()
+                exception = task.exception()
+                if exception:
+                    raise exception
 
+                workflow_result = task.result()
                 for index, workflow_step in enumerate(workflow_steps):
                     workflow_step_result = workflow_result.get("steps")[index]
 
@@ -284,11 +287,7 @@ async def send_message() -> AsyncIterable[str]:
                                 )
 
             except Exception as error:
-                yield (
-                    f"id: {workflow_step['agent_name']}\n"
-                    f"event: error\n"
-                    f"data: {error}\n\n"
-                )
+                yield (f"event: error\n" f"data: {error}\n\n")
 
                 if SEGMENT_WRITE_KEY:
                     for workflow_step in workflow_data.steps:

diff --git a/libs/superagent/app/utils/callbacks.py b/libs/superagent/app/utils/callbacks.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 from typing import Any, AsyncIterator, List, Literal, Tuple, Union, cast
 
 from decouple import config
@@ -9,6 +10,8 @@
 from langfuse import Langfuse
 from litellm import cost_per_token, token_counter
 
+logger = logging.getLogger(__name__)
+
 
 class CustomAsyncIteratorCallbackHandler(AsyncCallbackHandler):
     """Callback handler that returns an async iterator."""
@@ -17,6 +20,8 @@ class CustomAsyncIteratorCallbackHandler(AsyncCallbackHandler):
 
     done: asyncio.Event
 
+    TIMEOUT_SECONDS = 30
+
     @property
     def always_verbose(self) -> bool:
         return True
@@ -65,7 +70,13 @@ async def aiter(self) -> AsyncIterator[str]:
                     asyncio.ensure_future(self.done.wait()),
                 ],
                 return_when=asyncio.FIRST_COMPLETED,
+                timeout=self.TIMEOUT_SECONDS,
             )
+            # if we the timeout has been reached
+            if not done or not other:
+                logger.warning(f"{self.TIMEOUT_SECONDS} seconds of timeout reached")
+                self.done.set()
+                break
 
             # Cancel the other task
             if other:

diff --git a/libs/ui/app/workflows/[id]/chat.tsx b/libs/ui/app/workflows/[id]/chat.tsx
@@ -46,7 +46,12 @@ export default function Chat({
   )
   const [isLoading, setIsLoading] = React.useState<boolean>(false)
   const [messages, setMessages] = React.useState<
-    { type: string; message: string; steps?: Record<string, string> }[]
+    {
+      type: string
+      message: string
+      steps?: Record<string, string>
+      isSuccess?: boolean
+    }[]
   >(
     workflowSteps[0]?.agent?.initialMessage
       ? [{ type: "ai", message: workflowSteps[0].agent.initialMessage }]
@@ -147,13 +152,21 @@ export default function Chat({
                   type: "function_call",
                 },
               ])
-            }
+            } else if (event.event === "error") {
+              setMessages((previousMessages) => {
+                let updatedMessages = [...previousMessages]
 
-            if (
-              event.data !== "[END]" &&
-              event.event !== "function_call" &&
-              currentEventId
-            ) {
+                for (let i = updatedMessages.length - 1; i >= 0; i--) {
+                  if (updatedMessages[i].type === "ai") {
+                    updatedMessages[i].message = event.data
+                    updatedMessages[i].isSuccess = false
+                    break
+                  }
+                }
+
+                return updatedMessages
+              })
+            } else if (event.data !== "[END]" && currentEventId) {
               if (!messageByEventIds[currentEventId])
                 messageByEventIds[currentEventId] = ""
 
@@ -232,13 +245,14 @@ export default function Chat({
       <div className="relative flex flex-1 flex-col border-l bg-background text-sm">
         <ScrollArea className="flex-1 overflow-y-auto">
           <div className="flex flex-1 flex-col space-y-0 px-4 pb-10 pt-8">
-            {messages.map(({ type, message, steps }, index) => (
+            {messages.map(({ type, message, steps, isSuccess }, index) => (
               <Message
                 key={index}
                 type={type}
                 message={message}
                 steps={steps}
                 profile={profile}
+                isSuccess={isSuccess}
               />
             ))}
             <div ref={endOfMessagesRef} className="pt-20" />

diff --git a/libs/ui/components/message.tsx b/libs/ui/components/message.tsx
@@ -233,7 +233,7 @@ function MessageAlert({ error }: MessageAlertProps) {
       <AiOutlineExclamationCircle className="h-4 w-4" />
       <AlertTitle>Error</AlertTitle>
       <AlertDescription>
-        <b>{error}.</b>
+        <b>{error}</b>
       </AlertDescription>
     </Alert>
   )