Add OpenAI chatbot

.gitignore

@@ -1,4 +1,5 @@
 __pycache__
 .venv
 venv
-.vscode
+.vscode
+.env

autohdl/data.py

@@ -1,6 +1,19 @@
 import re
 from datasets import load_dataset
 
+
+def extract_header(prompt):
+    """Use regex to extract module header from prompt"""
+    
+    module_re = re.compile(r"module.*;")
+
+    try: 
+        result = re.search(module_re, prompt).group(0)
+    except:
+        raise Exception("Prompt is not in expected format when extracting header")
+    
+    return result.strip()
+
 def extract_description(prompt):
     """Use regex to extract description from prompt
     

autohdl/linter.py

@@ -0,0 +1,23 @@
+import subprocess
+
+def linter(code:str)-> str:
+    """A tool that runs Verilog syntax linting (with Verilator) on input code string
+    Args:
+        code: A string representing the Verilog code to be linted
+    """
+    
+    with open("code.v", "w") as f:
+        f.write(code)
+    
+    try:
+        subprocess.run(["verilator", "--lint-only", "-Wall", "code.v"],
+                                stdout=subprocess.PIPE,
+                                stderr=subprocess.STDOUT,
+                                check=True,
+                                encoding="utf-8")
+    
+        return "The linting passed successfully - no changes necessary"
+    
+    except subprocess.CalledProcessError as e:
+        return f"Verilator linting gave an error. Please investigate and fix: {e.stdout}"
+        

autohdl/llm.py

@@ -1,15 +1,62 @@
 import outlines
 from outlines.inputs import Chat
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from .data import data
+import openai
+import tiktoken
+from dotenv import load_dotenv
 
-system_prompt = ("You only complete chats with syntax correct Verilog code. "
-                 "End the Verilog module code completion with 'endmodule'. "
-                 "Do not include module, input and output definitions.")
 
 class LLM:
+    system_prompt = ("You only complete chats with syntax correct Verilog code. "
+                    "End the Verilog module code completion with 'endmodule'. "
+                    "Do not include module, input and output definitions.")
+    
     def __init__(self):
-        self.system_prompt = system_prompt
+        load_dotenv()
+
+
+class OpenAI(LLM):
+    def __init__(self, system_prompt=None, max_context_len=1000, model="gpt-5-nano"):
+        super().__init__()
+        
+        if system_prompt:
+            self.system_prompt = system_prompt
+            
+        self.max_context_len = max_context_len
+        self.model = model
+        
+        self.messages = [{"role": "system", "content": self.system_prompt}]
+        
+        self.client = openai.OpenAI()
+        
+    def __call__(self, message):
+        self.messages.append({"role": "user", "content": message})
+        
+        response = self.client.responses.create(
+            model=self.model,
+            input=self.messages,
+            service_tier="flex"
+        )
+        
+        self.messages.append({"role": "assistant", "content": response.output_text})
+        
+        return response.output_text
+
+    def truncate(self):
+        """Truncate off tokens until we reach max_context_len
+        
+        Approach is to first determine the index where we are below the max token length
+        in the messages array. Once we find this index, then we truncate off that content until
+        we are below the max.
+        """
+        total_len = sum(len(message['content']) for message in self.messages)
+        
+        i = len(self.messages)-1
+        while total_len > self.max_context_len:
+            pass
+        
+
+class HuggingFace(LLM):
         
     def load_model(self, use_cpu=True):
         if use_cpu:

autohdl/test_llm.py

@@ -0,0 +1,50 @@
+from .llm import OpenAI
+from .linter import linter
+from .data import data, extract_header
+
+
+def main():
+    llm = OpenAI()
+    
+    ds = data()
+    
+    message = ds['description'][0]["high_level_global_summary"]
+    expected = ds['code'][0]
+    
+    print("-------------------------------------------")
+    print("System prompt:")
+    print(llm.messages[0]['content'])
+    print("-------------------------------------------\n")
+    
+    print("-------------------------------------------")
+    print("User prompt:\n", message)
+    print("-------------------------------------------\n")
+    
+    print("-------------------------------------------")
+    print("Expected:\n", expected)
+    print("-------------------------------------------\n")
+    
+    response = llm(message)
+    print("-------------------------------------------")
+    print("Response:\n", response)
+    print("-------------------------------------------\n")
+    
+    header = extract_header(message)
+    lint_input = header + "\n" + response
+    print("-------------------------------------------")
+    print("Linter input:\n", lint_input)
+    print("-------------------------------------------\n") 
+    
+    lint_result = linter(lint_input)
+    print("-------------------------------------------")
+    print("Linter result:\n", lint_result)
+    print("-------------------------------------------\n")
+
+    second_response = llm(lint_result)
+    print("-------------------------------------------")
+    print("Second response:\n", second_response)
+    print("-------------------------------------------\n")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -2,4 +2,6 @@ transformers
 datasets
 outlines[transformers]
 streamlit
-streamlit_code_editor
+streamlit_code_editor
+python-dotenv
+openai