Merge pull request #18 from WhatTheFuzz/feature/background_thread

Implement query submission on BackgroundTaskThread

Author: WhatTheFuzz

README.md

@@ -82,12 +82,7 @@ for each model, as described in [OpenAI's documentation][tokens].
 
 ## Known Issues
 
-The query does not use Python's [asyncio][asyncio] and thus blocks the main
-thread. You may be unable to interact with the Binary Ninja UI while the query
-is waiting to be resolved. In some cases, your operating system may detect that
-Binary Ninja has stopped responding and ask you to Force Quit it. I have not
-experience any egregiously long hangs, however. This is documented in issue
-[#8][issue-8].
+Please submit an issue if you find something that isn't working properly.
 
 ## License
 
@@ -97,6 +92,4 @@ This project is licensed under the [MIT license][license].
 [token]:https://beta.openai.com/account/api-keys
 [tokens]:https://beta.openai.com/docs/models/gpt-3
 [entry]:./src/entry.py
-[asyncio]:https://docs.python.org/3/library/asyncio.html
-[issue-8]:https://github.com/WhatTheFuzz/binaryninja-openai/issues/8
 [license]:./LICENSE

plugin.json

@@ -27,6 +27,6 @@
          "openai"
       ]
 	},
-   "version": "1.2.0",
+   "version": "1.3.0",
    "minimumbinaryninjaversion": 3200
 }

src/agent.py

@@ -12,6 +12,8 @@
 from binaryninja.settings import Settings
 from binaryninja import log
 
+from . query import Query
+
 
 class Agent:
 
@@ -142,11 +144,9 @@ def generate_query(self, function: Union[LowLevelILFunction,
         prompt += '\n'.join(self.instruction_list(function))
         return prompt
 
-    def send_query(self, query: str) -> str:
-        '''Sends a query to the engine and returns the response.'''
-        response: str = openai.Completion.create(
-            model=self.model,
-            prompt=query,
-            max_tokens=self.get_token_count(),
-        )
-        return response.choices[0].text
+    def send_query(self, query: str) -> None:
+        '''Sends a query to the engine and prints the response.'''
+        query = Query(query_string=query,
+                      model=self.model,
+                      max_token_count=self.get_token_count())
+        query.start()

src/query.py

@@ -0,0 +1,25 @@
+import openai
+from binaryninja.plugin import BackgroundTaskThread
+
+
+class Query(BackgroundTaskThread):
+
+    def __init__(self, query_string: str, model: str,
+                 max_token_count: int) -> None:
+        BackgroundTaskThread.__init__(self,
+                                      initial_progress_text="",
+                                      can_cancel=False)
+        self.query_string: str = query_string
+        self.model: str = model
+        self.max_token_count: int = max_token_count
+
+    def run(self) -> None:
+        self.progress = "Submitting query to OpenAI."
+
+        response: str = openai.Completion.create(
+            model=self.model,
+            prompt=self.query_string,
+            max_tokens=self.max_token_count,
+        )
+        # Notify the user.
+        print(response.choices[0].text)