Fix chat output

.gitignore

@@ -164,3 +164,4 @@ cython_debug/
 
 # Other
 tempDir/picture.png
+

app.py

@@ -51,7 +51,7 @@ if picture is not None:
     with open(os.path.join("tempDir","picture.png"),"wb") as f:
          f.write(picture.getbuffer())
     img_url = f"https://huggingface.co/spaces/butterswords/MM_Math_Helper/resolve/main/tempDir/picture.png"
-    mmChat(img_url)
-    #guidedMM(st.session_state.systemPrompt, "http://192.168.50.36:8501/tempDir/picture")
+    #mmChat(img_url)
+    guidedMM(st.session_state.systemPrompt, img_url)
 else:
     basicChat()

helpers/chat.py

@@ -17,31 +17,65 @@ def hf_stream(model_name: str, messages: dict):
     for chunk in stream:
         chunk.choices[0].delta.content, end=""
 
-def hf_generator(model,prompt,data):
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "text",
-                    "text": prompt
-                },
-                {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": data
+def hf_generator(model,prompt,data,system=None):
+    if system:    
+        messages = [
+            {
+                "role": "system",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": system
                     }
-                }
-            ]
-        }
-    ]
+                ]
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": prompt
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": data
+                        }
+                    }
+                ]
+            }
+        ]
+    else:
+                messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": prompt
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": data
+                        }
+                    }
+                ]
+            }
+        ]
 
     completion = client.chat.completions.create(
         model=model, 
         messages=messages, 
         max_tokens=500
     )
-    return completion.choices[0].message
+    response = completion.choices[0].message.content
+    logger.info({"role": "assistant", "content": response})
+    st.session_state.messages.append({"role": "assistant", "content": response})
+    return completion.choices[0].message.content
+    
+
+    
 
 def basicChat():
 # Accept user input and then writes the response
@@ -80,7 +114,7 @@ def mmChat(data):
 
         with st.chat_message(st.session_state.model):
             logger.info(f"Message to {st.session_state.model}: {st.session_state.messages[-1]}")
-            response = st.write_stream(hf_generator(
+            response = st.write(hf_generator(
                 st.session_state.model,
                 prompt,
                 data))
@@ -96,34 +130,17 @@ def guidedMM(sysChoice:str, data):
     
     if prompt := st.chat_input("How may I help you learn math today?"):
          # Add user message to chat history
-        st.session_state.messages.append([
-	{
-		"role": "user",
-		"content": [
-			{
-				"type": "text",
-				"text": prompt
-			},
-			{
-				"type": "image_url",
-				"image_url": {
-					"url": data
-				}
-			}
-		]
-	}
-])
-        logger.info(st.session_state.messages[-2:])
+        st.session_state.messages.append({"role": "user", "content": prompt,"images":[data]})
         # Display user message in chat message container
         with st.chat_message("user"):
             st.markdown(prompt)
 
         with st.chat_message(st.session_state.model):
             logger.info(f"Message to {st.session_state.model}: {st.session_state.messages[-1]}")
-            response = st.write_stream(hf_generator(
+            response = st.write(hf_generator(
                 st.session_state.model,
-                [st.session_state.messages[-1]]
+                prompt,
+                data,
+                system
             ))
-            st.session_state.messages.append({"role": "assistant", "content": response})
-            logger.info(st.session_state.messages[-1])