wip

Author: kharvd

Dockerfile

@@ -2,10 +2,12 @@ FROM arm64v8/python:3.9.17-alpine
 WORKDIR /app
 
 RUN apk add build-base linux-headers clang rust cargo
-COPY . .
+COPY requirements.txt requirements_docker.txt ./
 RUN --mount=type=cache,target=/root/.cache/pip pip install -r requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip pip install -r requirements_docker.txt
 
+COPY . .
+
 RUN adduser -D gpt
 USER gpt
 RUN mkdir -p $HOME/.config/gpt-cli

gpt.py

@@ -29,6 +29,7 @@
     choose_config_file,
     read_yaml_config,
 )
+from gptcli.interpreter import CodeInterpreterListener
 from gptcli.llama import init_llama_models
 from gptcli.logging import LoggingChatListener
 from gptcli.cost import PriceChatListener
@@ -216,27 +217,6 @@ def run_non_interactive(args, assistant):
     simple_response(assistant, "\n".join(args.prompt), stream=not args.no_stream)
 
 
-def evaluate_python(source: str) -> str:
-    if not os.environ.get("GPTCLI_ALLOW_CODE_EXECUTION"):
-        return "unable to evaluate"
-
-    source = "%colors NoColor\n" + source
-
-    logger.info("Evaluating python code: %s", source)
-
-    try:
-        result = subprocess.check_output(
-            ["ipython", "-c", source], stderr=subprocess.STDOUT
-        ).decode("utf-8")
-    except subprocess.CalledProcessError as e:
-        result = e.output.decode("utf-8")
-    except Exception as e:
-        return traceback.format_exc()
-
-    logger.info("Result: %s", result)
-    return result
-
-
 class CLIChatSession(ChatSession):
     def __init__(self, assistant: Assistant, markdown: bool, show_price: bool):
         listeners = [
@@ -247,13 +227,13 @@ def __init__(self, assistant: Assistant, markdown: bool, show_price: bool):
         if show_price:
             listeners.append(PriceChatListener(assistant))
 
+        if os.environ.get("GPTCLI_ALLOW_CODE_EXECUTION") == "1":
+            listeners.append(CodeInterpreterListener("python_eval"))
+
         listener = CompositeChatListener(listeners)
         super().__init__(
             assistant,
             listener,
-            functions={
-                "python_eval": evaluate_python,
-            },
         )
 
 

gptcli/assistant.py

@@ -4,7 +4,7 @@
 import platform
 from typing import Any, Dict, Iterator, Optional, TypedDict, List
 
-from gptcli.completion import CompletionProvider, ModelOverrides, Message
+from gptcli.completion import Completion, CompletionProvider, ModelOverrides, Message
 from gptcli.google import GoogleCompletionProvider
 from gptcli.llama import LLaMACompletionProvider
 from gptcli.openai import OpenAICompletionProvider
@@ -16,12 +16,14 @@ class AssistantConfig(TypedDict, total=False):
     model: str
     temperature: float
     top_p: float
+    enable_code_execution: bool
 
 
 CONFIG_DEFAULTS = {
     "model": "gpt-3.5-turbo",
     "temperature": 0.7,
     "top_p": 1.0,
+    "enable_code_execution": False,
 }
 
 DEFAULT_ASSISTANTS: Dict[str, AssistantConfig] = {
@@ -89,7 +91,7 @@ def init_messages(self) -> List[Message]:
         return self.config.get("messages", [])[:]
 
     def supported_overrides(self) -> List[str]:
-        return ["model", "temperature", "top_p"]
+        return ["model", "temperature", "top_p", "enable_code_execution"]
 
     def _param(self, param: str, override_params: ModelOverrides) -> Any:
         # If the param is in the override_params, use that value
@@ -101,9 +103,15 @@ def _param(self, param: str, override_params: ModelOverrides) -> Any:
 
     def complete_chat(
         self, messages, override_params: ModelOverrides = {}, stream: bool = True
-    ) -> Iterator[str]:
+    ) -> Iterator[Completion]:
         model = self._param("model", override_params)
         completion_provider = get_completion_provider(model)
+
+        enable_code_execution = (
+            bool(self._param("enable_code_execution", override_params))
+            and os.environ.get("GPTCLI_ALLOW_CODE_EXECUTION") == "1"
+        )
+
         return completion_provider.complete(
             messages,
             {
@@ -112,6 +120,7 @@ def complete_chat(
                 "top_p": float(self._param("top_p", override_params)),
             },
             stream,
+            enable_code_execution,
         )
 
 

gptcli/cli.py

@@ -1,6 +1,8 @@
+import base64
 import logging
 import re
 import json
+from imgcat import imgcat
 from prompt_toolkit import PromptSession
 from prompt_toolkit.history import FileHistory
 from openai import OpenAIError, InvalidRequestError
@@ -76,27 +78,20 @@ def _format_function_call(self, function_call: FunctionCall) -> str:
     def print(self, message_delta: Message):
         self.current_message = merge_dicts(self.current_message, message_delta)
 
-        style = "green" if self.current_message.get("role") == "assistant" else "yellow"
-
         if self.markdown:
             assert self.live
-            text = ""
-            if self.current_message.get("role") == "function":
-                result = self.current_message.get("content") or ""
-                text = f"```\n# Function result:\n{result}\n```"
-            else:
-                text = self.current_message.get("content", "")
+            text = self.current_message.get("content", "")
 
             function_call = self.current_message.get("function_call")
             if function_call:
                 text += self._format_function_call(function_call)
 
-            content = Markdown(text, style=style)
+            content = Markdown(text, style="green")
             self.live.update(content)
             self.live.refresh()
         else:
             self.console.print(
-                Text(message_delta.get("content", ""), style=style), end=""
+                Text(message_delta.get("content", ""), style="green"), end=""
             )
 
     def __exit__(self, *args):
@@ -119,6 +114,22 @@ def __enter__(self):
     def on_message_delta(self, message_delta: Message):
         self.printer.print(message_delta)
 
+    def on_function_result(self, result: dict):
+        self.console.print(Text("Function result:", style="yellow"))
+        if "image/png" in result:
+            image_base64 = result["image/png"]
+            image_bytes = base64.b64decode(image_base64)
+            imgcat(image_bytes)
+        if "text/plain" in result:
+            text = result["text/plain"]
+            if self.markdown:
+                content = Markdown(
+                    f"```\n{text}\n```",
+                )
+            else:
+                content = Text(text, style="yellow")
+            self.console.print(content)
+
     def __exit__(self, *args):
         self.printer.__exit__(*args)
 

gptcli/completion.py

@@ -10,7 +10,7 @@ class FunctionCall(TypedDict, total=False):
 
 class Message(TypedDict, total=False):
     role: Required[str]
-    content: str
+    content: Optional[str]
     name: str
     function_call: FunctionCall
 
@@ -32,6 +32,7 @@ class ModelOverrides(TypedDict, total=False):
     model: str
     temperature: float
     top_p: float
+    enable_code_execution: bool
 
 
 class CompletionDelta(TypedDict):
@@ -47,6 +48,10 @@ class Completion(TypedDict):
 class CompletionProvider:
     @abstractmethod
     def complete(
-        self, messages: List[Message], args: dict, stream: bool = False
+        self,
+        messages: List[Message],
+        args: dict,
+        stream: bool = False,
+        enable_code_execution: bool = False,
     ) -> Iterator[Completion]:
         pass

gptcli/composite.py

@@ -1,8 +1,8 @@
-from gptcli.completion import FunctionCall, Message, ModelOverrides
+from gptcli.completion import Message, ModelOverrides
 from gptcli.session import ChatListener, ResponseStreamer
 
 
-from typing import List
+from typing import List, Optional
 
 
 class CompositeResponseStreamer(ResponseStreamer):
@@ -18,6 +18,10 @@ def on_message_delta(self, message_delta: Message):
         for streamer in self.streamers:
             streamer.on_message_delta(message_delta)
 
+    def on_function_result(self, result: dict):
+        for streamer in self.streamers:
+            streamer.on_function_result(result)
+
     def __exit__(self, *args):
         for streamer in self.streamers:
             streamer.__exit__(*args)
@@ -57,3 +61,9 @@ def on_chat_response(
     ):
         for listener in self.listeners:
             listener.on_chat_response(messages, response, overrides)
+
+    def on_function_call(self, function_name: str, **kwargs) -> Optional[str]:
+        for listener in self.listeners:
+            result = listener.on_function_call(function_name, **kwargs)
+            if result is not None:
+                return result

gptcli/openai.py

@@ -5,33 +5,56 @@
 
 from gptcli.completion import Completion, CompletionProvider, Message
 
+FUNCTIONS_SCHEMA = [
+    {
+        "name": "python_eval",
+        "description": "Evaluate an arbitrary Python snippet",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "source": {
+                    "type": "string",
+                    "description": "The Python code to evaluate",
+                },
+            },
+            "required": ["source"],
+        },
+    },
+    {
+        "name": "pip_install",
+        "description": "Install a Python package. The kernel will be restarted automatically after the package is installed.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "package": {
+                    "type": "string",
+                    "description": "The package to install",
+                },
+            },
+            "required": ["package"],
+        },
+    },
+]
+
 
 class OpenAICompletionProvider(CompletionProvider):
     def complete(
-        self, messages: List[Message], args: dict, stream: bool = False
+        self,
+        messages: List[Message],
+        args: dict,
+        stream: bool = False,
+        enable_code_execution: bool = False,
     ) -> Iterator[Completion]:
         kwargs = {}
         if "temperature" in args:
             kwargs["temperature"] = args["temperature"]
         if "top_p" in args:
             kwargs["top_p"] = args["top_p"]
 
-        functions = [
-            {
-                "name": "python_eval",
-                "description": "Evaluate an arbitrary Python snippet",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "source": {
-                            "type": "string",
-                            "description": "The Python code to evaluate",
-                        },
-                    },
-                    "required": ["source"],
-                },
-            },
-        ]
+        if enable_code_execution:
+            functions = FUNCTIONS_SCHEMA
+        else:
+            functions = []
 
         response_iter = cast(
             Any,