Add function execution

Author: gnachman

Makefile

@@ -1,7 +1,10 @@
 .PHONY: test upload
 
-test:
-	rm -rf testenv/ ; python3 -m venv testenv ; source testenv/bin/activate.csh ; pip install -e . ; testenv/bin/gptline
+run:
+	source testenv/bin/activate ; pip install -e . ; testenv/bin/gptline
+
+rebuild:
+	rm -rf testenv/ ; python3 -m venv testenv ; source testenv/bin/activate ; pip install -e . ; testenv/bin/gptline
 
 upload:
 	rm -f dist/*

requirements.txt

@@ -1,2 +1,3 @@
 openai
 prompt-toolkit
+simplejson

src/chat.py

@@ -1,38 +1,100 @@
-import sys
-import time
+import simplejson
+import enum
+import inspect
 import openai
+import sys
 import threading
+import time
+import typing
+from src.spin import spin
 
-def create_chat_with_spinner(messages, temperature):
-    chats = []
+def get_json_type_name(value):
+    if isinstance(value, str):
+        return "string"
+    elif isinstance(value, int):
+        return "integer"
+    elif isinstance(value, float):
+        return "number"
+    elif isinstance(value, bool):
+        return "boolean"
+    else:
+        return "null"
 
-    def show_spinner():
-        spinner = ["|", "/", "-", "\\"]
-        i = 0
-        while len(chats) == 0:
-            sys.stdout.write("\r" + spinner[i % 4])
-            sys.stdout.flush()
-            time.sleep(0.1)
-            i += 1
+def _json_schema(func):
+    api_info = {}
 
+    # Get function name
+    api_info["name"] = func.__name__
 
-    def create_chat_model():
-        chat = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo", messages=messages, stream=True, temperature=temperature
-        )
-        chats.append(chat)
+    # Get function description from docstring
+    docstring = inspect.getdoc(func)
+    if docstring:
+        api_info["description"] = docstring.split("\n\n")[0]
+
+    # Get function parameters
+    parameters = {}
+    parameters["type"] = "object"
+
+    properties = {}
+
+    signature = inspect.signature(func)
+    required = []
+    for param_name, param in signature.parameters.items():
+        param_info = {}
 
-    thread = threading.Thread(target=create_chat_model)
-    thread.start()
+        # Get parameter type from type hints
+        if typing.get_origin(param.annotation) is typing.Union and type(None) in typing.get_args(param.annotation):
+            # Handle Optional case
+            inner_type = typing.get_args(param.annotation)[0]
+            if issubclass(inner_type, enum.Enum):
+                param_info["type"] = get_json_type_name(inner_type.__members__[list(inner_type.__members__)[0]].value)
+                param_info["enum"] = [member.value for member in inner_type]
+            else:
+                param_info["type"] = get_json_type_name(inner_type.__name__)
+        elif issubclass(param.annotation, enum.Enum):
+            param_info["type"] = get_json_type_name(param.annotation.__members__[list(param.annotation.__members__)[0]].value)
+            param_info["enum"] = [member.value for member in param.annotation]
+        else:
+            param_info["type"] = get_json_type_name(param.annotation.__name__)
 
-    spinner_thread = threading.Thread(target=show_spinner)
-    spinner_thread.start()
+        # Get parameter description from docstring
+        if docstring and param_name in docstring:
+            param_info["description"] = docstring.split(param_name + ":")[1].split("\n")[0].strip()
 
-    thread.join()
-    spinner_thread.join()
+        # Check if parameter is required
+        if param.default == inspect.Parameter.empty:
+            required.append(param_name)
 
-    sys.stdout.write("\r \r")
-    return chats[0]
+        # Add parameter info to parameters dict
+        properties[param_name] = param_info
+
+    # Add parameters to api_info
+    parameters["properties"] = properties
+    parameters["required"] = required
+    api_info["parameters"] = parameters
+
+    return api_info
+
+def create_chat_with_spinner(messages, temperature, functions):
+    return create_chat(messages, temperature, functions, True)
+
+def create_chat(messages, temperature, functions, spinner=False):
+    def create_chat_model():
+        args = {
+            "model": "gpt-4",
+            "messages": messages,
+            "stream": True,
+            "temperature": temperature,
+        }
+        if functions:
+            args['functions'] = list(map(lambda f: _json_schema(f), functions))
+        return openai.ChatCompletion.create(**args)
+
+    if not spinner:
+        chats = []
+        return create_chat_model()
+    else:
+        return spin(create_chat_model)
 
 
 def suggest_name(chat_id, message):
@@ -48,3 +110,16 @@ def suggest_name(chat_id, message):
     name = chat_completion_resp.choices[0].message.content
     return (chat_id, name)
 
+def invoke(functions, name, args_str):
+    try:
+        args = simplejson.loads(args_str, strict=False)
+    except Exception as e:
+        print("")
+        print(f"ChatGPT called {name} with bad input: {args_str}")
+        raise ValueError("The function arguments did not form a valid JSON document")
+
+    for func in functions:
+        if func.__name__ == name:
+            return func(**args)
+    else:
+        raise ValueError(f"Function '{name}' not found.")

src/db.py

@@ -1,4 +1,5 @@
 import sqlite3
+from typing import Optional
 import fcntl
 import os
 import sys
@@ -45,9 +46,11 @@ def create_schema(self):
             id INTEGER PRIMARY KEY AUTOINCREMENT,
             chat_id INTEGER,
             role TEXT,
-            content TEXT,
+            content TEXT NULL,
             time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
             deleted INTEGER DEFAULT 0,
+            function_call_name TEXT,
+            function_call_arguments TEXT,
             FOREIGN KEY (chat_id) REFERENCES chats (id)
         )
         """
@@ -70,16 +73,17 @@ def create_chat(self, name=None):
         self.conn.commit()
         return chat_id
 
-    def add_message(self, chat_id: int, role: str, content: str):
+    def add_message(self, chat_id: int, role: str, content: str, function_call_name: Optional[str], function_call_arguments: Optional[str]):
         cursor = self.conn.cursor()  # Create a cursor
-        query = "INSERT INTO messages (chat_id, role, content) VALUES (?, ?, ?)"
-        cursor.execute(query, (chat_id, role, content))
+        query = "INSERT INTO messages (chat_id, role, content, function_call_name, function_call_arguments) VALUES (?, ?, ?, ?, ?)"
+        cursor.execute(query, (chat_id, role, content, function_call_name, function_call_arguments))
         self.conn.commit()
         last_message_id = cursor.lastrowid
 
-        fts_query = "INSERT INTO messages_fts (message_id, content) VALUES (?, ?)"
-        self.conn.execute(fts_query, (last_message_id, content.lower()))
-        self.conn.commit()
+        if content is not None and role != "function":
+            fts_query = "INSERT INTO messages_fts (message_id, content) VALUES (?, ?)"
+            self.conn.execute(fts_query, (last_message_id, content.lower()))
+            self.conn.commit()
 
         query = f"UPDATE chats SET last_update = CURRENT_TIMESTAMP WHERE id = ?"
         cursor.execute(query, (chat_id, ))
@@ -100,7 +104,7 @@ def get_message_by_id(self, message_id: int):
             raise IndexError("Index out of range")
 
     def get_message_by_index(self, chat_id: int, index: int):
-        query = "SELECT role, content, time, id, deleted FROM messages WHERE chat_id = ? ORDER BY id LIMIT 1 OFFSET ?"
+        query = "SELECT role, content, time, id, deleted, function_call_name, function_call_arguments FROM messages WHERE chat_id = ? ORDER BY id LIMIT 1 OFFSET ?"
         result = self.conn.execute(query, (chat_id, index)).fetchone()
         if result:
             return result

src/input_reader.py

@@ -21,6 +21,7 @@ class UserInput:
     query: Optional[str] = None
     regenerate = False
     edit = False
+    allow_execution = False
 
 @dataclass
 class Chat:
@@ -30,8 +31,9 @@ class Chat:
 
 
 # Returns UserInput
-def read_input(chats, current_chat_name, can_regen, placeholder):
+def read_input(chats, current_chat_name, can_regen, placeholder, allow_execution):
     result = UserInput()
+    result.allow_execution = allow_execution
 
     session = PromptSession()
     kb = KeyBindings()
@@ -55,6 +57,7 @@ def _(event):
     SEARCH = "$$$SEARCH"
     REGENERATE = "$$$REGENERATE"
     EDIT = "$$$EDIT"
+    TOGGLE_SETTING = "$$$TOGGLE_SETTING"
 
     @kb.add(Keys.F2)
     def _(event):
@@ -81,6 +84,15 @@ def _(event):
             app = get_app()
             app.exit(result=EDIT)
 
+    @kb.add(Keys.F7)
+    def _(event):
+        result.allow_execution = not result.allow_execution
+        if result.allow_execution:
+            print("Command execution enabled.")
+        else:
+            print("Command execution disabled.")
+        app = get_app()
+        app.exit(result=TOGGLE_SETTING)
     def read_search_query():
         search_session = PromptSession()
         return search_session.prompt(HTML("<b>Search: </b>"))
@@ -148,6 +160,10 @@ def bottom_toolbar():
                 if can_regen:
                     text += "  <b>F5</b>: Regenerate"
                     text += "  <b>F6</b>: Edit Last"
+                if result.allow_execution:
+                    text += "  <b>F7</b>: Disable Execution"
+                else:
+                    text += "  <b>F7</b>: Enable Execution"
                 return HTML(text)
             value = session.prompt(
                     "",
@@ -166,6 +182,8 @@ def bottom_toolbar():
             elif value == REGENERATE:
                 result.regenerate = True
                 return result
+            elif value == TOGGLE_SETTING:
+              continue
             elif value == EDIT:
                 result.edit = True
                 return result