Update app.py

app.py

@@ -11,6 +11,8 @@ import os
 from datetime import datetime
 from skyfield.api import load
 import matplotlib.pyplot as plt
+from io import BytesIO
+from PIL import Image
 
 # Load environment variables
 load_dotenv()
@@ -121,12 +123,13 @@ def process_query(query_text: str, vectorstore):
         prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
         prompt = prompt_template.format(context=context_text, question=query_text)
         
-        model = HuggingFaceHub(
-            repo_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
+        model = HuggingFaceEndpoint(
+            endpoint_url="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
             task="text2text-generation",
+        #    huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
             model_kwargs={"temperature": 0.5, "max_length": 512}
         )
-        response_text = model.predict(prompt)
+        response_text = model.invoke(prompt)
         
         sources = list(set([doc.metadata.get("source", "") for doc, _ in results]))
         return response_text, sources
@@ -297,13 +300,18 @@ def chat_interface(query_text):
             
             # Generate plot
             fig = plot_pladder(PLadder)
+            buf = BytesIO()
+            fig.savefig(buf, format='png')  # Save figure to buffer as PNG
+            buf.seek(0)
+            img = Image.open(buf)  # Convert to PIL image
+            plt.close(fig)  # Close the figure to free memory
             
             # Compile response text
             text = "Планетарная лестница: " + ", ".join(PLadder_ru) + "\n"
             text += "Размеры зон:\n" + "\n".join([f"Зона {i+1}: {size_str} {class_ru}" 
                                                   for i, (size_str, class_ru) in enumerate(ZSizes_ru)]) + "\n\n"
             text += "\n".join(responses)
-            return text, fig
+            return text, img
         
         else:
             # Handle regular RAG query