Spaces:
Runtime error
Runtime error
File size: 2,828 Bytes
2b0805d |
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 |
from typing import Optional
from langchain.embeddings import OpenAIEmbeddings
from langchain import LLMChain, PromptTemplate
from langchain.vectorstores import FAISS
from langchain.docstore import InMemoryDocstore
from src.baby_agi import BabyAGI
from langchain.agents import ZeroShotAgent, Tool
from langchain import OpenAI, SerpAPIWrapper, LLMChain
from constants import (
EMBEDDING_MODEL_NAME,
EMBEDDING_SIZE,
TODO_CHAIN_MODEL_NAME,
BABY_AGI_MODEL_NAME
)
def run_agent(
user_input,
num_iterations,
baby_agi_model=BABY_AGI_MODEL_NAME,
todo_chaining_model=TODO_CHAIN_MODEL_NAME,
embedding_model=EMBEDDING_MODEL_NAME
):
# Define your embedding model
embeddings_model = OpenAIEmbeddings(model=embedding_model)
# Initialize the vectorstore as empty
import faiss
embedding_size = EMBEDDING_SIZE
index = faiss.IndexFlatL2(embedding_size)
vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
todo_prompt = PromptTemplate.from_template(
"You are a planner who is an expert at coming up with a todo list for a given objective. Come up with a todo list for this objective: {objective}"
)
todo_chain = LLMChain(
llm=OpenAI(temperature=0, model_name=todo_chaining_model),
prompt=todo_prompt
)
search = SerpAPIWrapper()
tools = [
Tool(
name="Search",
func=search.run,
description="useful for when you need to answer questions about current events",
),
Tool(
name="TODO",
func=todo_chain.run,
description="useful for when you need to come up with todo lists. Input: an objective to create a todo list for. Output: a todo list for that objective. Please be very clear what the objective is!",
),
]
prefix = """You are an AI who performs one task based on the following objective: {objective}. Take into account these previously completed tasks: {context}."""
suffix = """Question: {task}
{agent_scratchpad}"""
prompt = ZeroShotAgent.create_prompt(
tools,
prefix=prefix,
suffix=suffix,
input_variables=["objective", "task", "context", "agent_scratchpad"],
)
OBJECTIVE = user_input
llm = OpenAI(temperature=0, model_name=baby_agi_model)
# Logging of LLMChains
verbose = False
# If None, will keep on going forever. Customize the number of loops you want it to go through.
max_iterations: Optional[int] = num_iterations
baby_agi = BabyAGI.from_llm(
prompt=prompt,
tools=tools,
llm=llm,
vectorstore=vectorstore,
verbose=verbose,
max_iterations=max_iterations
)
if (user_input):
baby_agi({"objective": OBJECTIVE})
|