Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 from smolagents.agents import ToolCallingAgent
-from smolagents import tool, LiteLLMModel
+from smolagents import tool, GeminiModel  # Assuming GeminiModel is the correct class
 from typing import Optional
 import cv2
 import pytesseract
@@ -8,8 +8,8 @@ from PIL import Image
 import io
 import numpy as np
 
-# Define the LiteLLMModel
-model = LiteLLMModel(model_id="gpt-4o")
+# Define the GeminiModel
+model = GeminiModel(api_key="AIzaSyCuK0Z50g5AGZrpvG54c8v_8rsOSRPjJJs", model_id="gemini-1.5-flash")
 
 @tool
 def extract_components(image_data: bytes) -> str:
@@ -54,8 +54,8 @@ uploaded_file = st.file_uploader("Upload a web design image", type=["png", "jpg"
 if st.button("Extract and Generate Code"):
     if uploaded_file is not None:
         image_data = uploaded_file.read()
-        components = agent.run(f"extract_components {image_data}")
-        code = agent.run(f"generate_code {components}")
+        components = agent.run("extract_components", image_data=image_data)
+        code = agent.run("generate_code", components=components)
         st.write("Extracted Components:", components)
         st.write("Generated Code:", code)
     else: