Merge pull request #185 from callanwu/master

Add support for replicate

Author: jianghuyihei

examples/Single_Agent/chat_bot/config.json

@@ -37,20 +37,8 @@
       ],
       "begin_role": "Yang bufan",
       "begin_query": "hello,What are you looking for me for?",
-      "LLM_type": "OpenAI",
-      "LLM": {
-        "temperature": 1.0,
-        "model": "gpt-3.5-turbo-16k-0613",
-        "log_path": "logs/Response_state"
-      },
       "agent_states": {
         "Yang bufan": {
-          "LLM_type": "OpenAI",
-          "LLM": {
-            "temperature": 1.0,
-            "model": "gpt-3.5-turbo-16k-0613",
-            "log_path": "logs/Yang_bufan"
-          },
           "style": {
             "role": "The director of a private detective agency, a cold detective"
           },

requirements.txt

@@ -3,7 +3,7 @@ fastapi
 google_api_python_client
 google_auth_oauthlib
 gradio
-langchain==0.0.329
+langchain
 numpy
 openai
 litellm
@@ -21,7 +21,7 @@ text2vec
 torch
 tqdm
 uvicorn
-pydantic==1.10.9
-typing_extensions==4.5.0
+pydantic
+typing_extensions
 serpapi
 google-search-results

src/agents/Agent/Agent.py

@@ -87,6 +87,10 @@ def from_config(cls, config_path):
                 current_state_begin_role = current_state["begin_role"] if "begin_role" in current_state else current_state["roles"][0]
                 agent_begins[state_name]["is_begin"] = current_state_begin_role==agent_role if "begin_role" in current_state else False
                 agent_begins[state_name]["begin_query"] = current_state["begin_query"] if "begin_query" in current_state else " "
+                if "LLM_type" not in current_state["agent_states"][agent_role]:
+                    current_state["agent_states"][agent_role]["LLM_type"] = config["LLM_type"]
+                if "LLM" not in current_state["agent_states"][agent_role]:
+                    current_state["agent_states"][agent_role]["LLM"] = config["LLM"]
                 agent_LLMs[state_name] = init_LLM("logs"+os.sep+f"{agent_name}",**current_state["agent_states"][agent_role])
             agents[agent_name] = cls(
                 agent_name,

src/agents/LLM/base_LLM.py

@@ -19,7 +19,6 @@ def __init__(self,**kwargs) -> None:
         super().__init__()
         self.MAX_CHAT_HISTORY = eval(
             os.environ["MAX_CHAT_HISTORY"]) if "MAX_CHAT_HISTORY" in os.environ else 10
-        
         self.model = kwargs["model"] if "model" in kwargs else "gpt-3.5-turbo-16k-0613"
         self.temperature = kwargs["temperature"] if "temperature" in  kwargs else 0.3
         self.log_path = kwargs["log_path"].replace("/",os.sep) if "log_path" in kwargs else "logs"
@@ -205,16 +204,15 @@ def get_response(self,
                         messages=messages,
                         functions=functions,
                         function_call=function_call,
-                        temperature=temperature,
-                        custom_llm_provider = "openai"
+                        temperature=temperature
                     )
                 else:
                     response = litellm.completion(
                         model=model,
                         messages=messages,
                         temperature=temperature,
-                        stream=stream,
-                        custom_llm_provider = "openai")
+                        stream=stream
+                    )
                 break
             except Exception as e:
                 print(e)
@@ -248,9 +246,9 @@ def init_LLM(default_log_path,**kwargs):
         )
     if LLM_type == "Replicate": 
         LLM = (
-            OpenAILLM(**kwargs["LLM"])
+            ReplicateLLM(**kwargs["LLM"])
             if "LLM" in kwargs
-            else OpenAILLM(model = "replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf",temperature=0.3,log_path=log_path)
+            else ReplicateLLM(model = "replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf",temperature=0.3,log_path=log_path)
         )
-        return LLM
+    return LLM
 

src/agents/utils.py

@@ -59,21 +59,20 @@ def get_embedding(sentence):
 
     if embed_model_name in ["text-embedding-ada-002"]:
         client = OpenAI(api_key=(
-        os.environ["EMBED_API_KEY"]
-        if "EMBED_API_KEY" in os.environ
-        else os.environ["API_KEY"]
-    ))
+            os.environ["EMBED_API_KEY"]
+            if "EMBED_API_KEY" in os.environ
+            else os.environ["API_KEY"]
+        ))
         if "PROXY" in os.environ:
             assert (
                 "http:" in os.environ["PROXY"] or "socks" in os.environ["PROXY"]
             ), "PROXY error,PROXY must be http or socks"
             client.proxies = {os.environ["PROXY"]}
         if "EMBED_API_BASE" in os.environ or "EMBED_BASE" in os.environ:
-            client.base_url = (
-        os.environ["EMBED_API_BASE"]
-        if "EMBED_API_BASE" in os.environ
-        else os.environ["API_BASE"]
-    )
+            client.base_url = (os.environ["EMBED_API_BASE"]
+                if "EMBED_API_BASE" in os.environ
+                else os.environ["API_BASE"]
+            )
         sentence = sentence.replace("\n", " ")
         embed = client.embeddings.create(input = sentence, model= embed_model_name,encoding_format="float").data[0].embedding
         embed = torch.tensor(embed, dtype=torch.float32)