Create main.py

main.py

@@ -0,0 +1,184 @@
+from langchain_openai import ChatOpenAI
+from langchain_core.prompts import ChatPromptTemplate
+import os
+from langchain_core.pydantic_v1 import BaseModel, Field
+import gradio as gr
+
+
+class code(BaseModel):
+    """Code output"""
+
+    prefix: str = Field(description="Description of the problem and approach")
+    imports: str = Field(description="Code block import statements")
+    code: str = Field(description="Code block not including import statements")
+    possible_errors: str = Field(description="Description of potential error and vulnerabilities in the code")
+    description = "Schema for code solutions to questions about Code."
+
+class revision(BaseModel):
+    """Revision output"""
+    
+    imports_revision: str = Field(description="Revision of imports")
+    code_revision: str = Field(description="Revision of code")
+    overall_evaluation: str = Field(description="Thorough evaluation of the imports and of the code")
+    description = "Schema for code solutions to questions about Code."
+
+def dict_to_string(d):
+    # Create a list of strings in the form '-key value'
+    parts = [f"{key}:\n\t{value}" for key, value in d.items()]
+    # Join the parts with a space separator
+    result = '\n'.join(parts)
+    return result
+
+def coder_reply_to_string(solution: code):
+    d = solution.__dict__
+    return dict_to_string(d)
+    
+class CodeGenerator:
+    def __init__(self, coder, revisor, maxiters):
+        self.coder = coder
+        self.revisor = revisor
+        self.maxiters = maxiters
+    def generate(self, prompt, context):
+        i = 0
+        while i < self.maxiters:
+            solution = self.coder.invoke({"context": context, "messages": [("user", prompt)]})
+            revision = self.revisor.invoke({"messages": [("user", coder_reply_to_string(solution))]})
+            context = coder_reply_to_string(revision)
+            i+=1
+        return solution
+    
+class modularized_code(BaseModel):
+    prefix: str = Field(description="Description of the problem and approach")
+    imports: str = Field(description="Code block import statements")
+    code: str = Field(description="Modularized code block not including import statements")
+    description = "Schema for code solutions to questions about Code."
+
+def check_code(solution):
+    try:
+        exec(solution.imports)
+    except Exception as e:
+        return f"An error occurred during import phase: {e}"
+    try:
+        exec(solution.imports+"\n"+solution.code)
+        return "Code ran succesfully"
+    except Exception as e:
+        return f"An error occurred during code execution phase: {e}"
+
+class CodeChecker:
+    def __init__(self, solution, checker):
+        self.solution = solution
+        self.checker = checker
+    def check_n_refactor(self, iterations):
+        for i in range(iterations):
+            try:
+                success_status = check_code(self.solution)
+            except Exception as e:
+                success_status = f"An error occurred while checking the code: {e}"
+            print("Success status: " + success_status)
+            if success_status == "Code run succesfully":
+                return self.solution
+            else:
+                self.solution = self.checker.invoke({"code": self.solution.imports + "\n" + self.solution.code, "errors": success_status})
+        return self.solution
+                
+# Grader prompt
+code_gen_prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """You are a coding assistant with expertise in python. Based on the following context (which can be code or revision suggestions):  \n ------- \n  {context} \n ------- \n Answer the user 
+    question. Ensure any code you provide can be executed \n 
+    with all required imports and variables defined. First, structure your answer with a description of the code solution. \n
+    Secondly list the imports. Thirdly list the functioning code block. Finally, describe potential errors one might encounter while executing the code. Here is the user question:""",
+        ),
+        ("placeholder", "{messages}"),
+    ]
+)
+
+revisor_prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """You are a coding revisor with expertise in python. From the prefix, import, code and potential errors produced by a programmer, you should provide a thorough review with:
+    - First what would you improve in the imports
+    - Second what would you improve in the code block
+    - Third an overall evaluation of the whole code solution produced by the programmer. Here is the programmer solution you should revise:""",
+        ),
+        ("placeholder", "{messages}"),
+    ]
+)
+
+modules_prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """Your job is to refactor code in a modularized way, and you have expertise in python. From the prefix, import, code and potential errors produced by a programmer, you should provide a refactored and modularized code with:
+    - A description of what you did
+    - The imports
+    - The refactored and modularized code. 
+    Orient code generation on this context (which may be user's request or some revision suggestions):
+    \n\n-----------\n{context}\n-------------\n\n
+    Here is the programmer solution you should refactor and modularize:""",
+        ),
+        ("placeholder", "{messages}"),
+    ]
+)
+
+checker_prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """Your job is to refactor code based on these errors:
+        
+            ----------
+            {errors}
+            ----------
+            
+            This is the code to refactor:
+            
+            ----------
+            {code}
+            ----------
+    You should then reply with:
+    - A description of what you did
+    - The imports
+    - The refactored code""",
+        ),
+    ]
+)
+
+import time
+
+def reply(message, history, api_key, context):
+    os.environ["OPENAI_API_KEY"] = api_key
+    expt_llm = "gpt-4o"
+    llm = ChatOpenAI(temperature=0, model=expt_llm)
+    code_gen_chain = code_gen_prompt | llm.with_structured_output(code)
+    code_revision_chain = revisor_prompt | llm.with_structured_output(revision)
+    aicodegenerator = CodeGenerator(code_gen_chain, code_revision_chain, 5)
+    solution = aicodegenerator.generate(prompt=message, context=context)
+    code_modules_chain = modules_prompt | llm.with_structured_output(modularized_code)
+    aicodemodularizer = CodeGenerator(code_modules_chain, code_revision_chain, 5)
+    modules_solution = aicodemodularizer.generate(prompt=coder_reply_to_string(solution), context=message)
+    code_checker_chain = checker_prompt | llm.with_structured_output(code)
+    aicodechecker = CodeChecker(modules_solution, code_checker_chain)
+    final_solution = aicodechecker.check_n_refactor(5)
+    response = "The final solution for your code is:\n\n```python" + final_solution.imports + "\n" + final_solution.code + "\n```"
+    this_hist = ""
+    for char in response:
+        this_hist+=char
+        time.sleep(0.001)
+        yield this_hist
+    
+api_key_user = gr.Textbox(label="OpenAI API key", type="password")
+context_user = gr.Textbox(label="Context", info="Add some contextual instructions for the model to know how to generate code", value="def hello_world():\n\tprint('Hello world!')\n\nhello_world()")
+
+chatbot = gr.Chatbot(height=400)
+additional_accordion = gr.Accordion(label="Parameters to be set before you start chatting", open=True)
+
+with gr.Blocks() as demo:
+    gr.HTML("<h1 align='center'>Self-Reviewing Coding Assistant🤖💻</h1>")
+    gr.Image('coding_assistant.png')
+    gr.ChatInterface(fn=reply, additional_inputs=[api_key_user, context_user], additional_inputs_accordion=additional_accordion, chatbot=chatbot)
+
+demo.launch(server_name="0.0.0.0", server_port=7860)