Akshayram1/smol_agent

1import streamlit as st2from smolagents.agents import ToolCallingAgent3from smolagents import tool, LiteLLMModel4from typing import Optional5import cv26import pytesseract7from PIL import Image8import io9import numpy as np10import base6411 12# Define the LiteLLMModel with OpenAI key13model = LiteLLMModel(model_id="gpt-4o", api_key="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")14 15@tool16def extract_components(image_data_base64: str) -> str:17    """18    Extract components from a web design image.19 20    Args:21        image_data_base64: The image data in base64 string format.22 23    Returns:24        A string describing the components found in the image.25    """26    image_data = base64.b64decode(image_data_base64)27    image = Image.open(io.BytesIO(image_data))28    gray = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2GRAY)29    components = pytesseract.image_to_string(gray)30    return components31 32@tool33def generate_code(components: str) -> str:34    """35    Generate code for the given components.36 37    Args:38        components: A string describing the components.39 40    Returns:41        The generated code for the components.42    """43    # This is a placeholder implementation. You can replace it with actual code generation logic.44    return f"Generated code for components: {components}"45 46# Define the ToolCallingAgent47agent = ToolCallingAgent(tools=[extract_components, generate_code], model=model)48 49# Streamlit app title50st.title("Web Design Component Extractor")51 52# File uploader for the web design image53uploaded_file = st.file_uploader("Upload a web design image", type=["png", "jpg", "jpeg"])54 55# Button to run the agent56if st.button("Extract and Generate Code"):57    if uploaded_file is not None:58        image_data = uploaded_file.read()59        image_data_base64 = base64.b64encode(image_data).decode('utf-8')60        components = agent.run(f"extract_components {image_data_base64}")61        code = agent.run(f"generate_code {components}")62        st.write("Extracted Components:", components)63        st.write("Generated Code:", code)64    else:65        st.write("Please upload an image.")