Make ActionRunner accept custom LLMs (#13)

autoagents/agents/search.py

@@ -14,11 +14,10 @@ from langchain.schema import AgentAction, AgentFinish
 from langchain.callbacks import get_openai_callback
 from langchain.callbacks.base import AsyncCallbackHandler
 from langchain.callbacks.manager import AsyncCallbackManager
-
+from langchain.base_language import BaseLanguageModel
 
 from autoagents.tools.tools import search_tool, note_tool, rewrite_search_query
 from autoagents.utils.logger import InteractionsLogger
-from autoagents.utils.utils import OpenAICred
 
 
 # Set up the base template
@@ -124,9 +123,8 @@ class CustomOutputParser(AgentOutputParser):
     class Config:
         arbitrary_types_allowed = True
     ialogger: InteractionsLogger
-    cred: OpenAICred
+    llm: BaseLanguageModel
     new_action_input: Optional[str]
-
     action_history = defaultdict(set)
 
     def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
@@ -154,7 +152,7 @@ class CustomOutputParser(AgentOutputParser):
         if action_input in self.action_history[action]:
             new_action_input = rewrite_search_query(action_input,
                                                     self.action_history[action],
-                                                    cred)
+                                                    self.llm)
             self.ialogger.add_message({"query_rewrite": True})
             self.new_action_input = new_action_input
             self.action_history[action].add(new_action_input)
@@ -168,8 +166,7 @@ class CustomOutputParser(AgentOutputParser):
 class ActionRunner:
     def __init__(self,
                  outputq,
-                 cred: OpenAICred,
-                 model_name: str,
+                 llm: BaseLanguageModel,
                  persist_logs: bool = False):
         self.ialogger = InteractionsLogger(name=f"{uuid.uuid4().hex[:6]}", persist=persist_logs)
         tools = [search_tool, note_tool]
@@ -179,7 +176,7 @@ class ActionRunner:
                 input_variables=["input", "intermediate_steps"],
                 ialogger=self.ialogger)
 
-        output_parser = CustomOutputParser(ialogger=self.ialogger, cred=cred)
+        output_parser = CustomOutputParser(ialogger=self.ialogger, llm=llm)
 
         class MyCustomHandler(AsyncCallbackHandler):
             def __init__(self):
@@ -225,10 +222,6 @@ class ActionRunner:
 
         handler = MyCustomHandler()
 
-        llm = ChatOpenAI(openai_api_key=cred.key,
-                         openai_organization=cred.org,
-                         temperature=0,
-                         model_name=model_name)
         llm_chain = LLMChain(llm=llm, prompt=prompt, callbacks=[handler])
         tool_names = [tool.name for tool in tools]
         for tool in tools:

autoagents/models/custom.py

@@ -0,0 +1,33 @@
+import requests
+
+from langchain.llms.base import LLM
+
+
+class CustomLLM(LLM):
+    @property
+    def _llm_type(self) -> str:
+        return "custom"
+
+    def _call(self, prompt: str, stop=None) -> str:
+        r = requests.post(
+            "http://localhost:8000/v1/chat/completions",
+            json={
+                "model": "283-vicuna-7b",
+                "messages": [{"role": "user", "content": prompt}],
+                "stop": stop
+            },
+        )
+        result = r.json()
+        return result["choices"][0]["message"]["content"]
+
+    async def _acall(self, prompt: str, stop=None) -> str:
+        r = requests.post(
+            "http://localhost:8000/v1/chat/completions",
+            json={
+                "model": "283-vicuna-7b",
+                "messages": [{"role": "user", "content": prompt}],
+                "stop": stop
+            },
+        )
+        result = r.json()
+        return result["choices"][0]["message"]["content"]

autoagents/spaces/app.py

@@ -9,7 +9,9 @@ import openai
 
 from autoagents.utils.constants import MAIN_HEADER, MAIN_CAPTION, SAMPLE_QUESTIONS
 from autoagents.agents.search import ActionRunner
-from autoagents.utils.utils import OpenAICred
+
+from langchain.chat_models import ChatOpenAI
+
 
 async def run():
     output_acc = ""
@@ -44,10 +46,9 @@ async def run():
 
         # Ask the user to enter their OpenAI API key
         if (api_key := st.sidebar.text_input("OpenAI api-key", type="password")):
-            cred = OpenAICred(api_key, None)
+            api_org = None
         else:
-            cred = OpenAICred(os.getenv("OPENAI_API_KEY"),
-                              os.getenv("OPENAI_API_ORG"))
+            api_key, api_org = os.getenv("OPENAI_API_KEY"), os.getenv("OPENAI_API_ORG")
         with st.sidebar:
             model_dict = {
                 "gpt-3.5-turbo": "GPT-3.5-turbo",
@@ -67,18 +68,22 @@ async def run():
             for q in SAMPLE_QUESTIONS:
                 st.markdown(f"*{q}*")
 
-        if not cred.key:
+        if not api_key:
             st.warning(
                 "API key required to try this app. The API key is not stored in any form. [This](https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key) might help."
             )
+        elif api_org and st.session_state.model_name == "gpt-4":
+            st.warning(
+                "The free API key does not support GPT-4. Please switch to GPT-3.5-turbo or input your own API key."
+            )
         else:
             outputq = asyncio.Queue()
-            runner = ActionRunner(
-                    outputq,
-                    cred=cred,
-                    model_name=st.session_state.model_name,
-                    persist_logs=True,
-                    )  # log to HF-dataset
+            runner = ActionRunner(outputq,
+                                  ChatOpenAI(openai_api_key=api_key,
+                                             openai_organization=api_org,
+                                             temperature=0,
+                                             model_name=st.session_state.model_name),
+                                  persist_logs=True)  # log to HF-dataset
 
             async def cleanup(e):
                 st.error(e)

autoagents/tools/tools.py

@@ -3,9 +3,7 @@ import os
 from duckpy import Client
 from langchain import PromptTemplate, OpenAI, LLMChain
 from langchain.agents import Tool
-from langchain.chat_models import ChatOpenAI
-
-from autoagents.utils.utils import OpenAICred
+from langchain.base_language import BaseLanguageModel
 
 
 MAX_SEARCH_RESULTS = 20  # Number of search results to observe at a time
@@ -55,15 +53,12 @@ note_tool = Tool(name="Notepad",
                    description=notepad_description)
 
 
-def rewrite_search_query(q: str, search_history, cred: OpenAICred) -> str:
+def rewrite_search_query(q: str, search_history, llm: BaseLanguageModel) -> str:
     history_string = '\n'.join(search_history)
     template ="""We are using the Search tool.
                  # Previous queries:
                  {history_string}. \n\n Rewrite query {action_input} to be
                  different from the previous ones."""
-    llm = ChatOpenAI(temperature=0,
-                     openai_api_key=cred.key,
-                     openai_organization=cred.org)
     prompt = PromptTemplate(template=template,
                             input_variables=["action_input", "history_string"])
     llm_chain = LLMChain(prompt=prompt, llm=llm)

autoagents/utils/utils.py

@@ -1,8 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional
-
-
-@dataclass
-class OpenAICred:
-    key: str
-    org: Optional[str]

test.py

@@ -1,17 +1,22 @@
 import os
 import asyncio
-from autoagents.agents.search import ActionRunner
-from langchain.callbacks import get_openai_callback
+
 from pprint import pprint
-import pdb
 from ast import literal_eval
 from multiprocessing import Pool, TimeoutError
 
+from autoagents.agents.search import ActionRunner
+from langchain.callbacks import get_openai_callback
+from langchain.chat_models import ChatOpenAI
+
+
 async def work(user_input):
     outputq = asyncio.Queue()
-
-    API_O = os.getenv("OPENAI_API_KEY")
-    runner = ActionRunner(outputq, api_key=API_O, model_name="gpt-3.5-turbo")
+    llm = ChatOpenAI(openai_api_key=os.getenv("OPENAI_API_KEY"),
+                     openai_organization=os.getenv("OPENAI_API_ORG"),
+                     temperature=0,
+                     model_name="gpt-3.5-turbo")
+    runner = ActionRunner(outputq, llm=llm)
     task = asyncio.create_task(runner.run(user_input, outputq))
 
     while True: