adding support for llama2 (replicate), anthropic and cohere

.python-version

@@ -0,0 +1 @@
+3.11.0

gpyt/__init__.py

@@ -4,21 +4,29 @@
 
 from gpyt.free_assistant import FreeAssistant
 from gpyt.palm_assistant import PalmAssistant
+from gpyt.lite_assistant import LiteAssistant
 
 from .app import gpyt
 from .args import USE_EXPERIMENTAL_FREE_MODEL
 from .assistant import Assistant
 from .config import MODEL, PROMPT
 
+
 # check for environment variable first
 API_KEY = os.getenv("OPENAI_API_KEY")
 PALM_API_KEY = os.getenv("PALM_API_KEY")
+ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
+COHERE_API_KEY = os.getenv("COHERE_API_KEY")
+REPLICATE_API_KEY = os.getenv("REPLICATE_API_KEY")
 DOTENV_PATH = os.path.expanduser("~/.env")
 
 if API_KEY is None or (isinstance(API_KEY, str) and not len(API_KEY)):
     result = dotenv_values(DOTENV_PATH)
     API_KEY = result.get("OPENAI_API_KEY", None)
     PALM_API_KEY = result.get("PALM_API_KEY", None)
+    ANTHROPIC_API_KEY = result.get("ANTHROPIC_API_KEY", None)
+    COHERE_API_KEY = result.get("COHERE_API_KEY", None)
+    REPLICATE_API_KEY = result.get("REPLICATE_API_KEY", None)
 
 
 assert (
@@ -46,10 +54,17 @@
 generate your own OpenAI API key (see above)
 
 """
-
+model_keys = {}
+if ANTHROPIC_API_KEY:
+    model_keys["ANTHROPIC_API_KEY"] = ANTHROPIC_API_KEY
+if COHERE_API_KEY:
+    model_keys["COHERE_API_KEY"] = COHERE_API_KEY
+if REPLICATE_API_KEY:
+    model_keys["REPLICATE_API_KEY"] = REPLICATE_API_KEY
 
 gpt = Assistant(api_key=API_KEY or "", model=MODEL, prompt=PROMPT)
 gpt4 = Assistant(api_key=API_KEY or "", model="gpt-4", prompt=PROMPT)
 free_gpt = FreeAssistant()
 palm = PalmAssistant(api_key=PALM_API_KEY)
-app = gpyt(assistant=gpt, free_assistant=free_gpt, palm=palm, gpt4=gpt4)
+litellm = LiteAssistant(model_keys=model_keys, model=MODEL, prompt=PROMPT)
+app = gpyt(assistant=gpt, free_assistant=free_gpt, palm=palm, gpt4=gpt4, lite_assistant = litellm)

gpyt/args.py

@@ -1,9 +1,11 @@
 import argparse
 
 parser = argparse.ArgumentParser()
+
 parser.add_argument(
     "--free", help="Use the gpt4free model. (experimental)", action="store_true"
 )
+
 parser.add_argument(
     "--palm", help="Use PaLM 2, Google's Premier LLM (new)", action="store_true"
 )
@@ -14,9 +16,17 @@
     action="store_true",
 )
 
+parser.add_argument(
+    "--litellm",
+    help="--litellm <model> . See https://github.com/BerriAI/litellm",
+)
+
 
 args = parser.parse_args()
 
 USE_EXPERIMENTAL_FREE_MODEL = args.free
 USE_PALM_MODEL = args.palm
 USE_GPT4 = args.gpt4
+USE_LITELLM = len(args.litellm) > 0 if hasattr(args, "litellm") and args.litellm is not None else False
+LITELLM_MODEL = args.litellm if USE_LITELLM else ""
+print(f"{USE_LITELLM=}, {LITELLM_MODEL=}")

gpyt/components/assistant_app.py

@@ -8,9 +8,16 @@
 from textual.widgets import Footer, Header, LoadingIndicator
 
 from gpyt.free_assistant import FreeAssistant
+from gpyt.lite_assistant import LiteAssistant
 from gpyt.palm_assistant import PalmAssistant
 
-from ..args import USE_EXPERIMENTAL_FREE_MODEL, USE_GPT4, USE_PALM_MODEL
+from ..args import (
+    USE_EXPERIMENTAL_FREE_MODEL,
+    USE_GPT4,
+    USE_PALM_MODEL,
+    USE_LITELLM,
+    LITELLM_MODEL,
+)
 from ..assistant import Assistant
 from ..conversation import Conversation, Message
 from ..id import get_id
@@ -49,30 +56,38 @@ def __init__(
         free_assistant: FreeAssistant,
         palm: PalmAssistant,
         gpt4: Assistant,
+        lite_assistant: LiteAssistant,
     ):
         super().__init__()
         self.assistant = assistant
         self._free_assistant = free_assistant
         self._palm = palm
         self._gpt4 = gpt4
+        self._lite_assistant = lite_assistant
         self.conversations: list[Conversation] = []
         self.active_conversation: Conversation | None = None
         self._convo_ids_added: set[str] = set()
         self.use_free_gpt = USE_EXPERIMENTAL_FREE_MODEL
         self.use_palm = USE_PALM_MODEL
         self.use_gpt4 = USE_GPT4
+        self.use_litellm = USE_LITELLM
         self.use_default_model = not (
-            self.use_free_gpt or self.use_palm or self.use_gpt4
+            self.use_free_gpt or self.use_palm or self.use_gpt4 or self.use_litellm
         )
         self.scrolled_during_response_stream = False
 
-    def _get_assistant(self) -> Assistant | FreeAssistant | PalmAssistant:
+    def _get_assistant(
+        self,
+    ) -> Assistant | FreeAssistant | PalmAssistant | LiteAssistant:
         if self.use_palm:
             return self._palm
         if self.use_free_gpt:
             return self._free_assistant
         if self.use_gpt4:
             return self._gpt4
+        if self.use_litellm:
+            self._lite_assistant.model = LITELLM_MODEL
+            return self._lite_assistant
 
         return self.assistant
 
@@ -85,6 +100,8 @@ def adjust_model_border_title(self) -> None:
             model = "PaLM 2 🌴"
         elif self.use_gpt4:
             model = "GPT 4"
+        elif self.use_litellm:
+            model = f"LITELLM w/ {LITELLM_MODEL}"
 
         self.user_input.border_title = f"Model: {model}"
 

gpyt/lite_assistant.py

@@ -0,0 +1,189 @@
+from typing import Generator
+
+import litellm 
+from litellm import completion
+import traceback
+import tiktoken
+import sys, os
+
+from .config import (
+    API_ERROR_FALLBACK,
+    SUMMARY_PROMPT,
+    APPROX_PROMPT_TOKEN_USAGE,
+    PRICING_LOOKUP,
+    MODEL_MAX_CONTEXT,
+)
+
+
+class LiteAssistant:
+    """
+    Represents an OpenAI GPT assistant.
+
+    It can generate responses based on user input and conversation history.
+    """
+
+    kDEFAULT_SUMMARY_FALLTHROUGH = "User Question"
+
+    def __init__(
+        self,
+        *,
+        model_keys: dict,
+        model: str,
+        prompt: str,
+        memory: bool = True,
+    ):
+        for key in model_keys:
+            if key == "ANTHROPIC_API_KEY":
+                litellm.anthropic_key = model_keys[key]
+            if key == "COHERE_API_KEY":
+                litellm.cohere_key = model_keys[key]
+            if key == "REPLICATE_API_KEY":
+                litellm.replicate_key = model_keys[key]
+        self.model = model
+        self.prompt = prompt
+        self.summary_prompt = SUMMARY_PROMPT
+        self.memory = memory
+        self.messages = [
+            {"role": "system", "content": self.prompt},
+        ]
+        self.error_fallback_message = API_ERROR_FALLBACK
+        self._encoding_engine = tiktoken.get_encoding(
+            "cl100k_base"
+        )  # gpt3.5/gpt4 encoding engine
+        self.input_tokens_this_convo = APPROX_PROMPT_TOKEN_USAGE
+        self.output_tokens_this_convo = 10
+        self.price_of_this_convo = self.get_default_price_of_prompt()
+
+    def set_history(self, new_history: list):
+        """Reassign all message history except for system message"""
+        sys_prompt = self.messages[0]
+        self.messages = new_history[::]
+        self.messages.insert(0, sys_prompt)
+
+    def get_tokens_used(self, message: str) -> int:
+        return len(self._encoding_engine.encode(message))
+
+    def get_default_price_of_prompt(self) -> float:
+        return self.get_approximate_price(
+            self.get_tokens_used(self.prompt), 0
+        ) + self.get_approximate_price(self.get_tokens_used(self.summary_prompt), 10)
+
+    def tokens_used_this_convo(self) -> int:
+        num = self.input_tokens_this_convo + self.output_tokens_this_convo
+        has_limit = MODEL_MAX_CONTEXT.get(self.model, None)
+        if not has_limit:
+            return 0
+
+        return num
+
+    def update_token_usage_for_input(self, in_message: str, out_message: str) -> None:
+        in_tokens_used = self.get_tokens_used(in_message)
+        out_tokens_used = self.get_tokens_used(out_message)
+        self.input_tokens_this_convo += self.input_tokens_this_convo + in_tokens_used
+        self.output_tokens_this_convo += self.output_tokens_this_convo + out_tokens_used
+        self.price_of_this_convo += self.get_approximate_price(
+            self.input_tokens_this_convo, self.output_tokens_this_convo
+        )
+
+    def get_approximate_price(self, _in: int | float, _out: int | float) -> float:
+        known_pricing = PRICING_LOOKUP.get(self.model, None)
+        if not known_pricing:
+            return 0.0
+        in_price_ratio = PRICING_LOOKUP[self.model][0]
+        out_price_ratio = PRICING_LOOKUP[self.model][1]
+
+        limit = MODEL_MAX_CONTEXT.get(self.model, None)
+        if not limit:
+            return 0
+
+        _in = min(_in, limit / 2)
+        _out = min(_out, limit / 2)
+
+        return (_in / 1000) * in_price_ratio + (_out / 1000) * out_price_ratio
+
+    def clear_history(self):
+        """
+        Wipe all message history during this conversation, except for the
+        first system message
+        """
+        self.messages = [self.messages[0]]
+        self.input_tokens_this_convo = APPROX_PROMPT_TOKEN_USAGE
+        self.output_tokens_this_convo = 10
+        self.price_of_this_convo = self.get_default_price_of_prompt()
+
+    def get_response_stream(self, user_input: str) -> Generator:
+        """
+        Use OpenAI API to retrieve a ChatCompletion response from a GPT model.
+
+        Memory can be configured so that the assistant forgets previous messages
+        you or it has sent. (saves tokens ($$$) as well)
+        """
+        if not self.memory:
+            self.clear_history()
+        self.messages.append({"role": "user", "content": user_input})
+
+        response = completion(  # type: ignore
+            model=self.model,
+            messages=self.messages,
+            stream=True,
+        )
+        # print(type(response))
+        if isinstance(response, dict):
+            # fake stream output if model doesn't support it
+            tmp_response = response
+            text_response = response['choices'][0]['message']['content']
+            tmp_response['choices'][0]['delta'] = {"content": ""}
+            for i in range(0, len(text_response), 8):
+                tmp_response['choices'][0]['delta']['content'] = text_response[i : i + 8]
+                yield tmp_response
+
+        # print(response["usage"])
+        return None
+
+        # return response  # type: ignore
+
+    def get_response(self, user_input: str) -> str:
+        """Get an entire string back from the assistant"""
+        messages = [
+            {"role": "system", "content": self.summary_prompt},
+            {"role": "user", "content": user_input.rstrip()},
+        ]
+        response = completion(model=self.model, messages=messages)
+
+        return response["choices"][0]["message"]["content"]  # type: ignore
+
+    def get_conversation_summary(self, initial_message: str) -> str:
+        """Generate a short 6 word or less summary of the user\'s first message"""
+        try:
+            summary = self.get_response(initial_message)
+        except:
+            return Assistant.kDEFAULT_SUMMARY_FALLTHROUGH
+
+        return summary
+
+    def log_assistant_response(self, final_response: str) -> None:
+        if not self.memory:
+            return
+
+        self.messages.append({"role": "assistant", "content": final_response})
+
+
+def _test() -> None:
+    from config import PROMPT
+    model_keys={"OPENAI_API_KEY": os.getenv("OPENAI_API_KEY")}
+    gpt = LiteAssistant(model_keys=model_keys, model="claude-instant-1", prompt=PROMPT)
+
+    # response = gpt.get_response("How do I make carrot cake?")
+    # summary = gpt.get_conversation_summary("How do I get into Dirt Biking?")
+    # print(f"{summary = }")
+    response_iterator = gpt.get_response_stream("hi how are you!")
+    for response in response_iterator:
+        try:
+            print(response["choices"][0]["delta"]["content"], end="", flush=True)
+        except:
+            traceback.print_exc()
+            continue
+
+
+if __name__ == "__main__":
+    _test()

gpyt/styles.cssx

@@ -109,10 +109,8 @@ Options.hidden {
   offset-x: -100%;
 }
 
-OptionCheckbox {
-  height: 1;
-}
-
-OptionCheckbox:focus-within {
-  background: $secondary-background-lighten-2;
+RadioButton {
+  padding: 0;
+  height: 3;
+  margin: 0;
 }

litellm_uuid.txt

@@ -0,0 +1 @@
+720218e6-e9f5-4c0d-9024-57ed6266b938