initial commit

.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+*.pyc
+*.pyo
+
+.env

requirements.txt

@@ -0,0 +1,4 @@
+gradio==4.39.0
+hatchet==1.4.0
+google-generativeai==0.7.2
+openai==1.37.0

src/models.py

@@ -0,0 +1,99 @@
+""" A light wrapper around a bunch of chat LLMs. The class should define a method that takes text input and returns a response from the model.
+"""
+from abc import ABC, abstractmethod
+from typing import Generator, Optional, AsyncGenerator
+import os
+import random
+import openai
+import google.generativeai as genai
+
+class ChatModel(ABC):
+    def __init__(self, name):
+        self.name = name
+
+    def __str__(self):
+        return self.name
+
+    def __repr__(self):
+        return self.name
+
+    @abstractmethod
+    def get_response(self, prompt) -> Generator[str, None, None]:
+        pass
+
+
+class DummyModel(ChatModel):
+
+    def __init__(self):
+        super().__init__("dummy")
+
+    def get_response(self, prompt: str) -> Generator[str, None, None]:
+        response = f"Dummy response to: {prompt}"
+        for idx in range(len(response)):
+            yield response[:idx+1]
+
+
+class OpenAIModel(ChatModel):
+
+    def __init__(self, model: str, client: openai.OpenAI):
+        super().__init__(model)
+        self.model = model
+        self.client = client
+
+    def get_response(self, prompt: str) -> Generator[str, None, None]:
+        stream = self.client.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "system", "content": "You are PerfGuru, a helpful assistant for assisting developers in identifying performance bottlenecks in their code and optimizing them."},
+                {"role": "user", "content": prompt}
+            ],
+            stream=True,
+            max_tokens=4096,
+        )
+        response = ""
+        for chunk in stream:
+            response += chunk.choices[0].delta.content or ""
+            yield response
+
+
+
+class GeminiModel(ChatModel):
+
+    def __init__(self, model: str, api_key: Optional[str] = None):
+        super().__init__(model)
+        if api_key:
+            genai.configure(api_key=api_key)
+
+        self.model = genai.GenerativeModel(model)
+        self.config = genai.types.GenerationConfig(
+            candidate_count=1,
+            max_output_tokens=4096,
+        )
+
+    def get_response(self, prompt: str) -> Generator[str, None, None]:
+        stream = self.model.generate_content(prompt, stream=True, generation_config=self.config)
+        response = ""
+        for chunk in stream:
+            response += chunk.text or ""
+            yield response
+
+
+AVAILABLE_MODELS = []
+
+#AVAILABLE_MODELS.append( DummyModel() )
+
+if os.environ.get("OPENAI_API_KEY"):
+    openai_client = openai.OpenAI()
+    AVAILABLE_MODELS.append( OpenAIModel("gpt-4o-mini", openai_client) )
+    AVAILABLE_MODELS.append( OpenAIModel("gpt-3.5-turbo", openai_client) )
+
+if os.environ.get("GOOGLE_API_KEY"):
+    AVAILABLE_MODELS.append( GeminiModel("gemini-1.5-flash") )
+    AVAILABLE_MODELS.append( GeminiModel("gemini-1.5-pro") )
+
+
+if not AVAILABLE_MODELS:
+    raise ValueError("No models available. Please set OPENAI_API_KEY or GOOGLE_API_KEY environment variables.")
+
+def select_random_model() -> ChatModel:
+    return random.choice(AVAILABLE_MODELS)

src/perfguru.py

@@ -0,0 +1,164 @@
+import gradio as gr
+import json
+import os
+import datetime
+from itertools import zip_longest
+
+from models import select_random_model
+from rag import select_random_formatter
+
+
+def error_helper(msg: str, duration: int = 10):
+    raise gr.Error(msg, duration=duration)
+
+
+def code_upload(code_file_select):
+    if code_file_select is None:
+        return gr.Button(interactive=False)
+    else:
+        return gr.Button(interactive=True)
+
+
+def chat_with_llms(prompt, code_files, profile_file, profile_type):
+    model1 = select_random_model()
+    model2 = select_random_model()
+    formatter1 = select_random_formatter()
+    formatter2 = select_random_formatter()
+
+    print(f"Selected models: {model1.name} and {model2.name}")
+
+    formatted1 = formatter1.format_prompt(prompt, code_files, profile_file, profile_type, error_fn=error_helper)
+    formatted2 = formatter2.format_prompt(prompt, code_files, profile_file, profile_type, error_fn=error_helper)
+
+    if formatted1 is None or formatted2 is None:
+        error_helper("Failed to format prompt. Please try again.")
+    
+    response1 = model1.get_response(formatted1)
+    response2 = model2.get_response(formatted2)
+
+    if response1 is None:
+        error_helper(f"Failed to get response from {model1.name}. Please try again.")
+
+    if response2 is None:
+        error_helper(f"Failed to get response from {model2.name}. Please try again.")
+
+    source1 = gr.Markdown(f"{model1.name} + {formatter1.name}", visible=False, elem_classes=["not-voted"])
+    source2 = gr.Markdown(f"{model2.name} + {formatter2.name}", visible=False, elem_classes=["not-voted"])
+
+    # set vote buttons to deactive
+    vote_buttons = gr.Button(interactive=False), gr.Button(interactive=False), gr.Button(interactive=False), gr.Button(interactive=False)
+
+    for c1, c2 in zip_longest(response1, response2):
+        yield c1 or gr.Textbox(), source1, formatted1, c2 or gr.Textbox(), source2, formatted2, *vote_buttons
+
+    vote_buttons = gr.Button(interactive=True), gr.Button(interactive=True), gr.Button(interactive=True), gr.Button(interactive=True)
+    yield c1 or gr.Textbox(), source1, formatted1, c2 or gr.Textbox(), source2, formatted2, *vote_buttons
+
+def log_interaction(prompt, vote, response1, model1, formatter1, full_prompt1, response2, model2, formatter2, full_prompt2):
+    interaction = {
+        "prompt": prompt,
+        "full_prompt1": full_prompt1,
+        "full_prompt2": full_prompt2,
+        "response1": response1,
+        "response2": response2,
+        "vote": vote,
+        "model1": model1,
+        "formatter1": formatter1,
+        "model2": model2,
+        "formatter2": formatter2,
+        "timestamp": datetime.datetime.now().isoformat()
+    }
+
+    log_file_path = "perf_guru_log.json"    
+    if os.path.exists(log_file_path):
+        with open(log_file_path, "r") as log_file:
+            logs = json.load(log_file)
+    else:
+        logs = []
+    
+    logs.append(interaction)
+    
+    # Write updated logs to file
+    with open(log_file_path, "w") as log_file:
+        json.dump(logs, log_file, indent=4)
+
+def handle_vote(prompt, vote, response1, source1, full_prompt1, response2, source2, full_prompt2):
+    model1, formatter1 = source1.split(" + ")
+    model2, formatter2 = source2.split(" + ")
+
+    label1_class = "voted" if vote == "Vote for Response 1" else "not-voted"
+    label2_class = "voted" if vote == "Vote for Response 2" else "not-voted"
+
+    log_interaction(prompt, vote, response1, model1, formatter1, full_prompt1, response2, model2, formatter2, full_prompt2)
+    return gr.Markdown(visible=True, elem_classes=[label1_class]), gr.Markdown(visible=True, elem_classes=[label2_class]), \
+        gr.Button(interactive=False), gr.Button(interactive=False), gr.Button(interactive=False), gr.Button(interactive=False)
+
+# Define the Gradio interface
+with gr.Blocks(css=".not-voted p { color: black; } .voted p { color: green; } .response { padding: 25px; } .response-md { padding: 20px; }") as interface:
+
+    gr.Markdown("""# PerfGuru: Code Performance Chatbot
+                
+    Welcome to PerfGuru! 
+                
+    This is a tool for assisting developers in identifying performance bottlenecks in their code and optimizing them using LLMs.
+    Upload your code files and a performance profile (if available) to get started. Then ask away!
+    This interface is primarily for data collecting and evaluation purposes. You will be presented outputs from two different LLMs and asked to vote on which response you find more helpful.
+                
+    ---""")
+
+    gr.Markdown("""## Upload Code Files and Performance Profile
+    
+                You must upload at least one source code file to proceed. You can also upload a performance profile if you have one.
+                Currently supported formats are HPCToolkit, CProfile, and Caliper.""")
+    with gr.Row():
+        code_files = gr.File(label="Upload Code File", file_count='multiple')
+
+        with gr.Column():
+            profile_type = gr.Dropdown(['No Profile', 'HPCToolkit', 'CProfile', "Caliper"], value='No Profile', multiselect=False, label="Select Profile Type")
+            profile_file = gr.File(label="Upload Performance Profile")
+    
+    gr.Markdown("---")
+    gr.Markdown("""## Ask a Question
+
+                Now you can ask a question about your code performance and chat with PerfGuru!
+                Once you receive two responses, vote on which one you found more helpful.""")
+    
+    default_question = "Can you help me identify and fix performance bugs in this code?"
+    prompt = gr.Textbox(label="Ask a question about your code performance", value=default_question)
+    
+    chat_button = gr.Button("Chat with PerfGuru", interactive=False)
+    with gr.Row(equal_height=True):
+        with gr.Column():
+            with gr.Accordion("Response 1", elem_classes=["response"]):
+                response1 = gr.Markdown(label="Response 1", visible=True, elem_classes=["response-md"])
+            source1 = gr.Markdown("", visible=False)
+            full_prompt1 = gr.Textbox("", visible=False)
+        with gr.Column():
+            with gr.Accordion("Response 2", elem_classes=["response"]):
+                response2 = gr.Markdown(label="Response 2", visible=True, elem_classes=["response-md"])
+            source2 = gr.Markdown("", visible=False)
+            full_prompt2 = gr.Textbox("", visible=False)
+
+    # use code_upload to toggle the status of the 'chat_button' based on whether a code file is uploaded or not
+    code_files.change(code_upload, inputs=[code_files], outputs=[chat_button])
+    
+    with gr.Row():
+        vote1_button = gr.Button("Vote for Response 1", interactive=False)
+        vote2_button = gr.Button("Vote for Response 2", interactive=False)
+        tie_button = gr.Button("Vote for Tie", interactive=False)
+        skip_button = gr.Button("Skip", interactive=False)
+    
+    vote_btns = [vote1_button, vote2_button, tie_button, skip_button]
+    for btn in vote_btns:
+        btn.click(handle_vote, inputs=[prompt, btn, response1, source1, full_prompt1, response2, source2, full_prompt2], outputs=[source1, source2, *vote_btns])
+
+    # final chat button
+    chat_button.click(
+        chat_with_llms, 
+        inputs=[prompt, code_files, profile_file, profile_type], 
+        outputs=[response1, source1, full_prompt1, response2, source2, full_prompt2, *vote_btns]
+    )
+
+# Launch the Gradio interface
+if __name__ == '__main__':
+    interface.launch()

src/profiles.py

@@ -0,0 +1,26 @@
+""" Helpers for loading performance profiles and extracting text from them.
+"""
+import json
+import os
+from typing import Literal, Optional
+
+import hatchet as ht
+
+
+class Profile:
+
+    def __init__(self, profile_path: os.PathLike, profile_type: Literal["HPCToolkit", "CProfile", "Caliper"]):
+        self.gf = self._load(profile_path, profile_type)
+
+    def _load(self, profile_path: os.PathLike, profile_type: Literal["HPCToolkit", "CProfile", "Caliper"]) -> ht.GraphFrame:
+        if profile_type == "HPCToolkit":
+            return ht.GraphFrame.from_hpctoolkit(profile_path)
+        elif profile_type == "CProfile":
+            return ht.GraphFrame.from_cprofile(profile_path)
+        elif profile_type == "Caliper":
+            return ht.GraphFrame.from_caliper(profile_path)
+        else:
+            raise ValueError(f"Profile type {profile_type} not supported.")
+        
+    def profile_to_tree_str(self) -> str:
+        return self.gf.tree()

src/rag.py

@@ -0,0 +1,70 @@
+""" Techniques for formatting the prompts that are passed to the LLMs.
+    These need to handle 2 major tasks:
+        1. Taking a directory of source code and embedding it in the prompt meaningfully (and possibly concatenating it).
+        2. Embedding a performance profile in the prompt if available.
+"""
+from abc import ABC, abstractmethod
+from typing import Optional, List, Mapping
+from os import PathLike
+from os.path import basename
+import random
+
+from profiles import Profile
+
+class PerfGuruPromptFormatter(ABC):
+
+    def __init__(self, name: str):
+        self.name = name
+
+    def _read_code_files(self, code_paths: List[PathLike]) -> Mapping[PathLike, str]:
+        code_files = {}
+        for code_path in code_paths:
+            with open(code_path, "r") as file:
+                code_files[code_path] = file.read()
+        return code_files
+    
+    def _read_profile(self, profile_path: PathLike, profile_type: str) -> Profile:
+        return Profile(profile_path, profile_type)
+
+    @abstractmethod
+    def format_prompt(self, prompt: str, code_paths: List[PathLike], profile_path: Optional[PathLike] = None, profile_type: Optional[str] = None, error_fn: Optional[callable] = None) -> str:
+        pass
+
+
+class BasicPromptFormatter(PerfGuruPromptFormatter):
+
+    def __init__(self):
+        super().__init__("basic")
+
+    def format_prompt(self, prompt: str, code_paths: List[PathLike], profile_path: Optional[PathLike] = None, profile_type: Optional[str] = None, error_fn: Optional[callable] = None) -> str:
+        if not code_paths:
+            if error_fn:
+                error_fn("No code files provided. At least one code file must be provided.")
+            return None
+        
+        concatenated_code = ""
+        code_file_contents = self._read_code_files(code_paths)
+        for code_path, content in code_file_contents.items():
+            fname = basename(code_path)
+            concatenated_code += f"{fname}:\n{content}\n\n"
+
+        if profile_path:
+            if not profile_type:
+                if error_fn:
+                    error_fn("Profile type must be provided if a profile file is provided.")
+                return None
+            
+            profile = self._read_profile(profile_path, profile_type)
+            profile_content = profile.profile_to_tree_str()
+        else:
+            profile_content = ""
+
+        return f"Code:\n{concatenated_code}\n\n{profile_type} Profile:\n{profile_content}\n\n{prompt}"
+    
+
+AVAILABLE_FORMATTERS = []
+AVAILABLE_FORMATTERS.append(BasicPromptFormatter())
+
+
+def select_random_formatter() -> PerfGuruPromptFormatter:
+    return random.choice(AVAILABLE_FORMATTERS)