Implement query submission in background thread.

Author: WhatTheFuzz

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)