o3-mini agent gives code successfully
This commit is contained in:
Madhu
parent
76715d6e82
commit
ee8166d746
2 changed files with 250 additions and 10 deletions
|
|
@ -1,20 +1,255 @@
|
|||
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
|
||||
load_dotenv()
|
||||
|
||||
vision_agent = Agent(
|
||||
model=Gemini(id="gemini-2.0-flash-exp", api_key=os.getenv("GEMINI_API_KEY")),
|
||||
markdown=True,
|
||||
)
|
||||
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
|
||||
|
||||
coding_agent = Agent(
|
||||
model=OpenAIChat(id="o3-mini", api_key=os.getenv("OPENAI_API_KEY")),
|
||||
markdown=True
|
||||
)
|
||||
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")
|
||||
|
||||
# Print the response in the terminal
|
||||
coding_agent.print_response("Share a 2 sentence horror story.")
|
||||
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-2.0-flash-exp", 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."""
|
||||
try:
|
||||
if st.session_state.sandbox:
|
||||
st.session_state.sandbox.close()
|
||||
os.environ['E2B_API_KEY'] = st.session_state.e2b_key
|
||||
st.session_state.sandbox = Sandbox()
|
||||
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]:
|
||||
"""
|
||||
Run code in e2b sandbox and return execution results.
|
||||
|
||||
Args:
|
||||
code: Python code to execute
|
||||
|
||||
Returns:
|
||||
Dict containing execution logs and any output
|
||||
"""
|
||||
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.
|
||||
|
||||
Args:
|
||||
vision_agent: Initialized Gemini vision agent
|
||||
image: Uploaded image to process
|
||||
|
||||
Returns:
|
||||
str: Extracted problem description in natural language
|
||||
"""
|
||||
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."""
|
||||
|
||||
# Convert image to bytes for Gemini
|
||||
img_byte_arr = BytesIO()
|
||||
image.save(img_byte_arr, format=image.format)
|
||||
img_byte_arr = img_byte_arr.getvalue()
|
||||
|
||||
response = vision_agent.run(prompt, images=[img_byte_arr])
|
||||
return response.content
|
||||
|
||||
def execute_code_with_agent(execution_agent: Agent, code: str, sandbox: Sandbox) -> str:
|
||||
"""
|
||||
Use execution agent to run and explain code results.
|
||||
"""
|
||||
try:
|
||||
execution = sandbox.run_code(code)
|
||||
|
||||
# Handle execution errors
|
||||
if execution.error:
|
||||
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}"
|
||||
|
||||
prompt = f"""Here is the code execution result:
|
||||
Logs: {execution.logs}
|
||||
Files: {sandbox.files.list("/")}
|
||||
|
||||
Please provide a clear explanation of the results and any outputs."""
|
||||
|
||||
response = execution_agent.run(prompt)
|
||||
return response.content
|
||||
except Exception as e:
|
||||
return f"⚠️ Sandbox Error: {str(e)}"
|
||||
|
||||
def main() -> None:
|
||||
"""Main application function."""
|
||||
st.title("O3-Mini Coding Assistant")
|
||||
|
||||
initialize_session_state()
|
||||
setup_sidebar()
|
||||
|
||||
# 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_column_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..."):
|
||||
image = Image.open(uploaded_image)
|
||||
extracted_query = process_image_with_gemini(vision_agent, image)
|
||||
|
||||
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)
|
||||
|
||||
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..."):
|
||||
if not st.session_state.sandbox:
|
||||
initialize_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)
|
||||
|
||||
# Display any generated files
|
||||
files = st.session_state.sandbox.files.list("/")
|
||||
if files:
|
||||
st.markdown("📁 **Generated Files:**")
|
||||
st.json(files)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
|
|
|||
5
ai_agent_tutorials/o3-mini-agent/requirements.txt
Normal file
5
ai_agent_tutorials/o3-mini-agent/requirements.txt
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
streamlit
|
||||
python-dotenv
|
||||
e2b
|
||||
agno
|
||||
Pillow
|
||||
Loading…
Reference in a new issue