Update app.py

app.py

@@ -567,10 +567,11 @@ def setup_document_chat(processed_docs):
                     ])
                     
                     # Process with model
-                    response = llm.invoke({
-                        "context": context,
-                        "question": query
-                    }) | llm
+                    messages = template.format_messages(
+                        context=context,
+                        question=query
+                    )
+                    response = llm.invoke(messages)
                     
                     return {"answer": response}
                 
@@ -851,7 +852,8 @@ with tab2:
                                         ]
                                         template = ChatPromptTemplate.from_messages(messages)
                                         
-                                        response = llm.invoke(template.invoke({"input": prompt}))
+                                        messages = template.format_messages(input=prompt)
+                                        response = llm.invoke(messages)
                                         
                                         # Try to parse the response into the pydantic model
                                         try: