final version
This commit is contained in:
Madhu
parent
d8825c00ee
commit
ceb218bc50
2 changed files with 0 additions and 299 deletions
|
|
@ -1,294 +0,0 @@
|
||||||
from typing import Optional, Dict, Any
|
|
||||||
import streamlit as st
|
|
||||||
from agno.agent import Agent, RunResponse
|
|
||||||
from agno.models.openai import OpenAIChat
|
|
||||||
from agno.models.google import Gemini
|
|
||||||
from e2b_code_interpreter import Sandbox
|
|
||||||
import os
|
|
||||||
from dotenv import load_dotenv
|
|
||||||
from PIL import Image
|
|
||||||
from io import BytesIO
|
|
||||||
import base64
|
|
||||||
load_dotenv()
|
|
||||||
|
|
||||||
def initialize_session_state() -> None:
|
|
||||||
"""Initialize Streamlit session state variables."""
|
|
||||||
if 'openai_key' not in st.session_state:
|
|
||||||
st.session_state.openai_key = ''
|
|
||||||
if 'gemini_key' not in st.session_state:
|
|
||||||
st.session_state.gemini_key = ''
|
|
||||||
if 'e2b_key' not in st.session_state:
|
|
||||||
st.session_state.e2b_key = ''
|
|
||||||
if 'sandbox' not in st.session_state:
|
|
||||||
st.session_state.sandbox = None
|
|
||||||
|
|
||||||
def setup_sidebar() -> None:
|
|
||||||
"""Setup sidebar with API key inputs."""
|
|
||||||
with st.sidebar:
|
|
||||||
st.title("API Configuration")
|
|
||||||
st.session_state.openai_key = st.text_input("OpenAI API Key",
|
|
||||||
value=st.session_state.openai_key,
|
|
||||||
type="password")
|
|
||||||
st.session_state.gemini_key = st.text_input("Gemini API Key",
|
|
||||||
value=st.session_state.gemini_key,
|
|
||||||
type="password")
|
|
||||||
st.session_state.e2b_key = st.text_input("E2B API Key",
|
|
||||||
value=st.session_state.e2b_key,
|
|
||||||
type="password")
|
|
||||||
|
|
||||||
def create_agents() -> tuple[Agent, Agent, Agent]:
|
|
||||||
"""Create vision, coding, and execution agents with API keys from session state."""
|
|
||||||
vision_agent = Agent(
|
|
||||||
model=Gemini(id="gemini-exp-1206", api_key=st.session_state.gemini_key),
|
|
||||||
markdown=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
coding_agent = Agent(
|
|
||||||
model=OpenAIChat(
|
|
||||||
id="o3-mini",
|
|
||||||
api_key=st.session_state.openai_key,
|
|
||||||
system_prompt="""You are an expert Python programmer. You will receive coding problems similar to LeetCode questions,
|
|
||||||
which may include problem statements, sample inputs, and examples. Your task is to:
|
|
||||||
1. Analyze the problem carefully
|
|
||||||
2. Write clean, efficient Python code to solve it
|
|
||||||
3. Include proper documentation and type hints
|
|
||||||
4. The code will be executed in an e2b sandbox environment
|
|
||||||
Please ensure your code is complete and handles edge cases appropriately."""
|
|
||||||
),
|
|
||||||
markdown=True
|
|
||||||
)
|
|
||||||
|
|
||||||
execution_agent = Agent(
|
|
||||||
model=OpenAIChat(
|
|
||||||
id="o3-mini",
|
|
||||||
api_key=st.session_state.openai_key,
|
|
||||||
system_prompt="""You are an expert at executing Python code in sandbox environments.
|
|
||||||
Your task is to:
|
|
||||||
1. Take the provided Python code
|
|
||||||
2. Execute it in the e2b sandbox
|
|
||||||
3. Format and explain the results clearly
|
|
||||||
4. Handle any execution errors gracefully
|
|
||||||
Always ensure proper error handling and clear output formatting."""
|
|
||||||
),
|
|
||||||
markdown=True
|
|
||||||
)
|
|
||||||
|
|
||||||
return vision_agent, coding_agent, execution_agent
|
|
||||||
|
|
||||||
def initialize_sandbox() -> None:
|
|
||||||
"""Initialize or reset the e2b sandbox with proper timeout configuration."""
|
|
||||||
try:
|
|
||||||
if st.session_state.sandbox:
|
|
||||||
try:
|
|
||||||
st.session_state.sandbox.close()
|
|
||||||
except:
|
|
||||||
pass
|
|
||||||
os.environ['E2B_API_KEY'] = st.session_state.e2b_key
|
|
||||||
# Initialize sandbox with 60 second timeout
|
|
||||||
st.session_state.sandbox = Sandbox(timeout=60)
|
|
||||||
except Exception as e:
|
|
||||||
st.error(f"Failed to initialize sandbox: {str(e)}")
|
|
||||||
st.session_state.sandbox = None
|
|
||||||
|
|
||||||
def run_code_in_sandbox(code: str) -> Dict[str, Any]:
|
|
||||||
if not st.session_state.sandbox:
|
|
||||||
initialize_sandbox()
|
|
||||||
|
|
||||||
execution = st.session_state.sandbox.run_code(code)
|
|
||||||
return {
|
|
||||||
"logs": execution.logs,
|
|
||||||
"files": st.session_state.sandbox.files.list("/")
|
|
||||||
}
|
|
||||||
|
|
||||||
def process_image_with_gemini(vision_agent: Agent, image: Image) -> str:
|
|
||||||
"""
|
|
||||||
Process uploaded image with Gemini Vision to extract code problem.
|
|
||||||
"""
|
|
||||||
prompt = """Analyze this image and extract any coding problem or code snippet shown.
|
|
||||||
Describe it in clear natural language, including any:
|
|
||||||
1. Problem statement
|
|
||||||
2. Input/output examples
|
|
||||||
3. Constraints or requirements
|
|
||||||
Format it as a proper coding problem description."""
|
|
||||||
|
|
||||||
# Save image to a temporary file
|
|
||||||
temp_path = "temp_image.png"
|
|
||||||
try:
|
|
||||||
# Convert to RGB if needed
|
|
||||||
if image.mode != 'RGB':
|
|
||||||
image = image.convert('RGB')
|
|
||||||
image.save(temp_path, format="PNG")
|
|
||||||
|
|
||||||
# Read the file and create image data
|
|
||||||
with open(temp_path, 'rb') as img_file:
|
|
||||||
img_bytes = img_file.read()
|
|
||||||
|
|
||||||
# Pass image to Gemini
|
|
||||||
response = vision_agent.run(
|
|
||||||
prompt,
|
|
||||||
images=[{"filepath": temp_path}] # Use filepath instead of content
|
|
||||||
)
|
|
||||||
return response.content
|
|
||||||
except Exception as e:
|
|
||||||
st.error(f"Error processing image: {str(e)}")
|
|
||||||
return "Failed to process the image. Please try again or use text input instead."
|
|
||||||
finally:
|
|
||||||
# Clean up temporary file
|
|
||||||
if os.path.exists(temp_path):
|
|
||||||
os.remove(temp_path)
|
|
||||||
|
|
||||||
def execute_code_with_agent(execution_agent: Agent, code: str, sandbox: Sandbox) -> str:
|
|
||||||
"""
|
|
||||||
Use execution agent to run and explain code results.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
# Set timeout to 30 seconds for code execution
|
|
||||||
sandbox.set_timeout(30)
|
|
||||||
execution = sandbox.run_code(code)
|
|
||||||
|
|
||||||
# Handle execution errors
|
|
||||||
if execution.error:
|
|
||||||
if "TimeoutException" in str(execution.error):
|
|
||||||
return "⚠️ Execution Timeout: The code took too long to execute (>30 seconds). Please optimize your solution or try a smaller input."
|
|
||||||
|
|
||||||
error_prompt = f"""The code execution resulted in an error:
|
|
||||||
Error: {execution.error}
|
|
||||||
|
|
||||||
Please analyze the error and provide a clear explanation of what went wrong."""
|
|
||||||
response = execution_agent.run(error_prompt)
|
|
||||||
return f"⚠️ Execution Error:\n{response.content}"
|
|
||||||
|
|
||||||
# Get files list safely
|
|
||||||
try:
|
|
||||||
files = sandbox.files.list("/")
|
|
||||||
except:
|
|
||||||
files = []
|
|
||||||
|
|
||||||
prompt = f"""Here is the code execution result:
|
|
||||||
Logs: {execution.logs}
|
|
||||||
Files: {str(files)}
|
|
||||||
|
|
||||||
Please provide a clear explanation of the results and any outputs."""
|
|
||||||
|
|
||||||
response = execution_agent.run(prompt)
|
|
||||||
return response.content
|
|
||||||
except Exception as e:
|
|
||||||
# Reinitialize sandbox on error
|
|
||||||
try:
|
|
||||||
initialize_sandbox()
|
|
||||||
except:
|
|
||||||
pass
|
|
||||||
return f"⚠️ Sandbox Error: {str(e)}"
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
"""Main application function."""
|
|
||||||
st.title("O3-Mini Coding Agent")
|
|
||||||
|
|
||||||
# Add timeout info in sidebar
|
|
||||||
initialize_session_state()
|
|
||||||
setup_sidebar()
|
|
||||||
with st.sidebar:
|
|
||||||
st.info("⏱️ Code execution timeout: 30 seconds")
|
|
||||||
|
|
||||||
# Check all required API keys
|
|
||||||
if not (st.session_state.openai_key and
|
|
||||||
st.session_state.gemini_key and
|
|
||||||
st.session_state.e2b_key):
|
|
||||||
st.warning("Please enter all required API keys in the sidebar.")
|
|
||||||
return
|
|
||||||
|
|
||||||
vision_agent, coding_agent, execution_agent = create_agents()
|
|
||||||
|
|
||||||
# Clean, single-column layout
|
|
||||||
uploaded_image = st.file_uploader(
|
|
||||||
"Upload an image of your coding problem (optional)",
|
|
||||||
type=['png', 'jpg', 'jpeg']
|
|
||||||
)
|
|
||||||
|
|
||||||
if uploaded_image:
|
|
||||||
st.image(uploaded_image, caption="Uploaded Image", use_container_width=True)
|
|
||||||
|
|
||||||
user_query = st.text_area(
|
|
||||||
"Or type your coding problem here:",
|
|
||||||
placeholder="Example: Write a function to find the sum of two numbers. Include sample input/output cases.",
|
|
||||||
height=100
|
|
||||||
)
|
|
||||||
|
|
||||||
# Process button
|
|
||||||
if st.button("Generate & Execute Solution", type="primary"):
|
|
||||||
if uploaded_image and not user_query:
|
|
||||||
# Process image with Gemini
|
|
||||||
with st.spinner("Processing image..."):
|
|
||||||
try:
|
|
||||||
# Save uploaded file to temporary location
|
|
||||||
image = Image.open(uploaded_image)
|
|
||||||
extracted_query = process_image_with_gemini(vision_agent, image)
|
|
||||||
|
|
||||||
if extracted_query.startswith("Failed to process"):
|
|
||||||
st.error(extracted_query)
|
|
||||||
return
|
|
||||||
|
|
||||||
st.info("📝 Extracted Problem:")
|
|
||||||
st.write(extracted_query)
|
|
||||||
|
|
||||||
# Pass extracted query to coding agent
|
|
||||||
with st.spinner("Generating solution..."):
|
|
||||||
response = coding_agent.run(extracted_query)
|
|
||||||
except Exception as e:
|
|
||||||
st.error(f"Error processing image: {str(e)}")
|
|
||||||
return
|
|
||||||
|
|
||||||
elif user_query and not uploaded_image:
|
|
||||||
# Direct text input processing
|
|
||||||
with st.spinner("Generating solution..."):
|
|
||||||
response = coding_agent.run(user_query)
|
|
||||||
|
|
||||||
elif user_query and uploaded_image:
|
|
||||||
st.error("Please use either image upload OR text input, not both.")
|
|
||||||
return
|
|
||||||
else:
|
|
||||||
st.warning("Please provide either an image or text description of your coding problem.")
|
|
||||||
return
|
|
||||||
|
|
||||||
# Display and execute solution
|
|
||||||
if 'response' in locals():
|
|
||||||
st.divider()
|
|
||||||
st.subheader("💻 Solution")
|
|
||||||
|
|
||||||
# Extract code from markdown response
|
|
||||||
code_blocks = response.content.split("```python")
|
|
||||||
if len(code_blocks) > 1:
|
|
||||||
code = code_blocks[1].split("```")[0].strip()
|
|
||||||
|
|
||||||
# Display the code
|
|
||||||
st.code(code, language="python")
|
|
||||||
|
|
||||||
# Execute code with execution agent
|
|
||||||
with st.spinner("Executing code..."):
|
|
||||||
# Always initialize a fresh sandbox for each execution
|
|
||||||
initialize_sandbox()
|
|
||||||
|
|
||||||
if st.session_state.sandbox:
|
|
||||||
execution_results = execute_code_with_agent(
|
|
||||||
execution_agent,
|
|
||||||
code,
|
|
||||||
st.session_state.sandbox
|
|
||||||
)
|
|
||||||
|
|
||||||
# Display execution results
|
|
||||||
st.divider()
|
|
||||||
st.subheader("🚀 Execution Results")
|
|
||||||
st.markdown(execution_results)
|
|
||||||
|
|
||||||
# Try to display files if available
|
|
||||||
try:
|
|
||||||
files = st.session_state.sandbox.files.list("/")
|
|
||||||
if files:
|
|
||||||
st.markdown("📁 **Generated Files:**")
|
|
||||||
st.json(files)
|
|
||||||
except:
|
|
||||||
pass
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
|
|
||||||
|
|
@ -1,5 +0,0 @@
|
||||||
streamlit
|
|
||||||
python-dotenv
|
|
||||||
e2b
|
|
||||||
agno
|
|
||||||
Pillow
|
|
||||||
Loading…
Reference in a new issue