arzzit/Python_Code_Generator
0
1# Python Code Assistant - Design Documentation2 3## Core Architecture4# The Python Code Assistant is built as a Streamlit web application that integrates Google's Gemini Pro model through the LangChain framework.5# The application follows a simple flow: users input their API key and code requirements, the system validates the key, 6# processes the request through the LLM, and displays generated code with optional test results. The core strength lies in its modular design, 7# separating concerns between UI handling, API authentication, code generation, and response processing.8 9## Design Choices and Implementation10# The implementation prioritizes reliability and user experience through several key features. 11# A retry mechanism handles temporary API failures, while comprehensive error handling ensures graceful degradation when issues occur. 12# The prompt template is structured to generate consistent, well-documented Python code following PEP 8 guidelines and includes proper error handling. 13# The response format is strictly defined using tags ([CODE] and [TEST RESULTS]) to ensure reliable parsing and display of results.14 15## Current Limitations and Assumptions16# The system operates under several practical assumptions: users have valid API keys and basic Python knowledge, 17# and moderate query complexity is expected. The current implementation handles single user sessions and 18# processes one request at a time. Memory management is basic, with simple cleanup after each response processing. 19# The application assumes reasonable response times from the API and doesn't currently implement caching or advanced optimization.20 21## Future Development Path22# The most impactful improvements would focus on three areas: enhanced functionality, reliability, and user experience. 23# Key additions could include support for multiple programming languages, interactive code editing, and result caching for similar queries. 24# The security layer could be strengthened with proper API key encryption and rate limiting. User experience could be improved with loading indicators, 25# syntax highlighting for test results, and one-click code copying functionality. 26# These improvements would maintain the application's simplicity while expanding its capabilities and reliability.27 28import streamlit as st29from langchain_google_genai import ChatGoogleGenerativeAI30from langchain.prompts import PromptTemplate31from langchain.chains import LLMChain32from langchain_community.chat_message_histories import ChatMessageHistory33from langchain_core.chat_history import BaseChatMessageHistory34import re35import time36 37# Validate the API key38def valid_api(apikey):39 try:40 llm = ChatGoogleGenerativeAI(model="gemini-1.5-pro", api_key=apikey)41 ans = llm.invoke("test input")42 return True if ans else False43 except Exception:44 return False45 46# Retry mechanism for LLM invocation47def invoke_with_retry(chain, session_id, query, testcase, retries=3, delay=2):48 for attempt in range(retries):49 try:50 input_data = {"query": query, "testcases": testcase}51 response = chain.run(input_data)52 return response53 except Exception as e:54 if attempt < retries - 1:55 time.sleep(delay) # Retry delay56 else:57 st.error(f"Agent failed after {retries} attempts: {e}")58 return None59 60#main function61def main():62 # Streamlit UI63 st.title("Python Code Assistant")64 api_key = st.text_input("Enter your Gemini API key", type="password") # Your API key65 66 if api_key: # Only proceed when API key is entered67 if valid_api(api_key): # Check if the entered API key is valid68 llm = ChatGoogleGenerativeAI(model="gemini-1.5-pro", api_key=api_key)69 70 # the prompt template71 prompt = PromptTemplate(72 input_variables=["query", "testcases"],73 template=(74 """ You are a Python programming expert. 75 Generate clean, efficient, and well-documented Python code based on the user's requirements.76 77 Requirements: {query}78 79 Please follow these guidelines:80 1. Write well-documented code with clear docstrings81 2. Include appropriate error handling82 3. Use type hints where relevant83 4. Follow PEP 8 style guidelines84 5. Handle edge cases85 86 {testcases}87 88 Format your response exactly as follows:89 90 [CODE]91 <Write your Python code here>92 [END CODE]93 94 [TEST RESULTS]95 <Show test results if {testcases} provided>96 <Return None if no {testcases} provided>97 <If {testcases} provided is invalid return Invalid testcase>98 [END TEST RESULTS]99 100 Important:101 - If test cases are provided, show for each test:102 * Input: <actual input>103 * Expected: <expected output>104 * Result: <actual output>105 * Status: PASS/FAIL106 - If no test cases are provided, only show the code section107 - Don't explain the code unless specifically asked108 - Don't show multiple solutions unless requested109 - Don't add any text outside the specified format110 """111 ),112 )113 114 # LLM chain115 chain = LLMChain(llm=llm, prompt=prompt)116 117 # Function to extract code blocks and test result from AI response118 def extract_code_and_tests(response: str):119 """120 Extract both code and test results from the AI response.121 122 Args:123 response (str): Raw response from the AI124 125 Returns:126 tuple: (code, test_results) where both are strings127 """128 # Extract code between [CODE] tags129 code_match = re.search(r'\[CODE\](.*?)\[END CODE\]', response, re.DOTALL)130 code = code_match.group(1).strip() if code_match else "No code found."131 132 # Extract test results between [TEST RESULTS] tags133 test_match = re.search(r'\[TEST RESULTS\](.*?)\[END TEST RESULTS\]', response, re.DOTALL)134 test_results = test_match.group(1).strip() if test_match else ""135 136 return code, test_results137 138 # Inputs139 st.header("Enter Your Query")140 user_input = st.text_input("Query", placeholder="e.g., Check if a string is a palindrome")141 user_testcase = st.text_input("Testcases (Optional)", placeholder="Sample valid testcases")142 143 if user_input:144 try:145 # Always provide a default value for `testcases`146 testcase_value = user_testcase if user_testcase else "No testcases provided"147 148 # Invoke chain and process response149 response = invoke_with_retry(150 chain=chain, 151 session_id="default_session", 152 query=user_input, 153 testcase=testcase_value154 )155 if response:156 code, test_results = extract_code_and_tests(response)157 158 # Display the code159 st.subheader("Generated Code:")160 st.code(code, language="python", line_numbers=True)161 162 # Display test results if they exist163 if test_results != 'None':164 st.subheader("Test Results:")165 st.write(test_results)166 else:167 st.error("Failed to process the query.")168 169 # release memory170 del response171 172 except KeyError as e:173 st.error(f"Error: Missing key in response - {e}")174 except Exception as e:175 st.error(f"Unexpected error: {str(e)}")176 else:177 st.error("Invalid API Key. Please try again.")178 else:179 st.info("Please enter your Gemini API key to start.")180 181if __name__ == "__main__":182 main()183 184 185 186 187 188 189 190 191 192 193