Spaces:
Runtime error
Runtime error
File size: 2,974 Bytes
b3509ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import os
import openai
import tiktoken
from swarmai.utils.ai_engines.EngineBase import EngineBase
from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.llms import OpenAI
from langchain.utilities import GoogleSearchAPIWrapper
class LanchainGoogleEngine(EngineBase):
"""
gpt-4, gpt-4-0314, gpt-4-32k, gpt-4-32k-0314, gpt-3.5-turbo, gpt-3.5-turbo-0301
"""
SUPPORTED_MODELS = [
"gpt-4",
"gpt-4-0314",
"gpt-4-32k",
"gpt-4-32k-0314",
"gpt-3.5-turbo",
"gpt-3.5-turbo-0301"
]
def __init__(self, model_name: str, temperature: float, max_response_tokens: int):
if model_name not in self.SUPPORTED_MODELS:
raise ValueError(f"Model {model_name} is not supported. Supported models are: {self.SUPPORTED_MODELS}")
super().__init__("openai", model_name, temperature, max_response_tokens)
if "OPENAI_API_KEY" not in os.environ:
raise ValueError("OPENAI_API_KEY environment variable is not set.")
openai.api_key = os.getenv("OPENAI_API_KEY")
self.tiktoken_encoding = tiktoken.encoding_for_model(model_name)
self.agent = self._init_chain()
self.search = GoogleSearchAPIWrapper()
def _init_chain(self):
"""Instantiates langchain chain with all the necessary tools
"""
llm = OpenAI(temperature=self.temperature)
tools = load_tools(["google-search", "google-search-results-json"], llm=llm)
agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=False, return_intermediate_steps=True)
return agent
def call_model(self, conversation: list) -> str:
"""Does the search itself but provides very short answers!
"""
if isinstance(conversation, list):
prompt = self._convert_conversation_to_str(conversation)
else:
prompt = conversation
response = self.agent(prompt)
final_response = ""
intermediate_steps = response["intermediate_steps"]
for step in intermediate_steps:
final_response += step[0].log + "\n" + step[1]
final_response += response["output"]
return final_response
def google_query(self, query: str) -> str:
"""Does the search itself but provides very short answers!
"""
response = self.search.run(query)
return response
def search_sources(self, query: str, n=5):
"""Does the search itself but provides very short answers!
"""
response = self.search.results(query, n)
return response
def _convert_conversation_to_str(self, conversation):
"""Converts conversation to a string
"""
prompt = ""
for message in conversation:
prompt += message["content"] + "\n"
return prompt
|