From 42f503cc1625eb7ba9f5f0860955710811203c2f Mon Sep 17 00:00:00 2001
From: Ben Lopata <benlopata@gmail.com>
Date: Fri, 18 Oct 2024 10:45:24 -0500
Subject: [PATCH] Wraps honcho operations in promise handling.

---
 api/routers/chat.py | 63 ++++++++++++++++++++++-----------------------
 1 file changed, 31 insertions(+), 32 deletions(-)

diff --git a/api/routers/chat.py b/api/routers/chat.py
index 6341db2..700e18b 100644
--- a/api/routers/chat.py
+++ b/api/routers/chat.py
@@ -12,12 +12,11 @@
 
 
 @router.post("/stream")
-async def stream(inp: schemas.ConversationInput, background_tasks: BackgroundTasks):
-    """Stream the response to the user, currently only used by the Web UI and has integration to be able to use Honcho if not anonymous"""
+async def stream(inp: schemas.ConversationInput):
     try:
         user = honcho.apps.users.get_or_create(app_id=app.id, name=inp.user_id)
 
-        def convo_turn():
+        async def convo_turn():
             thought = ""
             response = ""
             try:
@@ -50,8 +49,7 @@ def convo_turn():
                 yield f"Error: {str(e)}"
                 return
 
-            background_tasks.add_task(
-                create_messages_and_metamessages,
+            await create_messages_and_metamessages(
                 app.id, user.id, inp.conversation_id, inp.message, thought, response
             )
 
@@ -69,45 +67,24 @@ def convo_turn():
                 content={"error": "internal_server_error", "message": "An internal server error has occurred."}
             )
 
-
-@router.get("/thought/{message_id}")
-async def get_thought(conversation_id: str, message_id: str, user_id: str):
-    try:
-        user = honcho.apps.users.get_or_create(app_id=app.id, name=user_id)
-        thought = honcho.apps.users.sessions.metamessages.list(
-            session_id=conversation_id,
-            app_id=app.id,
-            user_id=user.id,
-            message_id=message_id,
-            metamessage_type="thought",
-        )
-        # In practice, there should only be one thought per message
-        return {"thought": thought.items[0].content if thought.items else None}
-    except Exception as e:
-        logging.error(f"An error occurred: {str(e)}")
-        return JSONResponse(
-            status_code=500,
-            content={"error": "internal_server_error", "message": "An internal server error has occurred."}
-        )
-
-
-def create_messages_and_metamessages(app_id, user_id, conversation_id, user_message, thought, ai_response):
+async def create_messages_and_metamessages(app_id, user_id, conversation_id, user_message, thought, ai_response):
     try:
-        honcho.apps.users.sessions.messages.create(
+        # These operations will use the DB layer's built-in retry logic
+        await honcho.apps.users.sessions.messages.create(
             is_user=True,
             session_id=str(conversation_id),
             app_id=app_id,
             user_id=user_id,
             content=user_message,
         )
-        new_ai_message = honcho.apps.users.sessions.messages.create(
+        new_ai_message = await honcho.apps.users.sessions.messages.create(
             is_user=False,
             session_id=str(conversation_id),
             app_id=app_id,
             user_id=user_id,
             content=ai_response,
         )
-        honcho.apps.users.sessions.metamessages.create(
+        await honcho.apps.users.sessions.metamessages.create(
             app_id=app_id,
             session_id=str(conversation_id),
             user_id=user_id,
@@ -116,4 +93,26 @@ def create_messages_and_metamessages(app_id, user_id, conversation_id, user_mess
             content=thought,
         )
     except Exception as e:
-        logging.error(f"Error in background task: {str(e)}")
+        logging.error(f"Error in create_messages_and_metamessages: {str(e)}")
+        raise  # Re-raise the exception to be handled by the caller
+
+
+@router.get("/thought/{message_id}")
+async def get_thought(conversation_id: str, message_id: str, user_id: str):
+    try:
+        user = honcho.apps.users.get_or_create(app_id=app.id, name=user_id)
+        thought = honcho.apps.users.sessions.metamessages.list(
+            session_id=conversation_id,
+            app_id=app.id,
+            user_id=user.id,
+            message_id=message_id,
+            metamessage_type="thought",
+        )
+        # In practice, there should only be one thought per message
+        return {"thought": thought.items[0].content if thought.items else None}
+    except Exception as e:
+        logging.error(f"An error occurred: {str(e)}")
+        return JSONResponse(
+            status_code=500,
+            content={"error": "internal_server_error", "message": "An internal server error has occurred."}
+        )