Refactor chat history structure in history.py and index.py

FloRul · Mar 21, 2024 · 01fc0d9 · 01fc0d9
1 parent 7847dcf
commit 01fc0d9
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 5 deletions.
diff --git a/lambdas/inference/src/history.py b/lambdas/inference/src/history.py
@@ -25,8 +25,28 @@ def get(self, limit: int = 10):
                 json.loads(response["Payload"].read().decode("utf-8-sig"))["body"]
             )
             for x in body:
-                result.append({"role": "user", "content": x["HumanMessage"]})
-                result.append({"role": "assistant", "content": x["AssistantMessage"]})
+                result.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": x["HumanMessage"],
+                            },
+                        ],
+                    }
+                )
+                result.append(
+                    {
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": x["AssistantMessage"],
+                            },
+                        ],
+                    }
+                )
             return result
         except ClientError as e:
             print("Error occurred: ", e.response["Error"]["Message"])

diff --git a/lambdas/inference/src/index.py b/lambdas/inference/src/index.py
@@ -114,9 +114,9 @@ def lambda_handler(event, context):
     try:
         source = event.get("queryStringParameters", {}).get("source", "message")
         embedding_collection_name = event["queryStringParameters"]["collectionName"]
-        
+
         logger.info(str(event))
-        
+
         sessionId = str(uuid.uuid1())
 
         if "sessionId" in event["queryStringParameters"]:
@@ -142,7 +142,17 @@ def lambda_handler(event, context):
 
         logger.info("fetching chat history...")
         chat_history = history.get(limit=5)
-        chat_history.append({"role": "user", "content": query})
+        chat_history.append(
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": query,
+                    },
+                ],
+            }
+        )
 
         logger.info(f"Chat history: {chat_history}")