add llm customization (run-llama#4)

1_🏠_Home.py

@@ -1,12 +1,8 @@
 import streamlit as st
-import os
-os.environ["OPENAI_API_KEY"] = st.secrets.openai_key
-
 from streamlit_pills import pills
 
 from agent_utils import (
     load_meta_agent_and_tools,
-    ParamCache
 )
 
 
@@ -15,7 +11,6 @@
 ####################
 
 
-
 st.set_page_config(page_title="Build a RAGs bot, powered by LlamaIndex", page_icon="🦙", layout="centered", initial_sidebar_state="auto", menu_items=None)
 st.title("Build a RAGs bot, powered by LlamaIndex 💬🦙")
 st.info(
@@ -24,6 +19,13 @@
     icon="ℹ️"
 )
 
+# TODO: noodle on this
+# with st.sidebar:
+#     openai_api_key_st = st.text_input("OpenAI API Key (optional, not needed if you filled in secrets.toml)", value="", type="password")
+#     if st.button("Save"):
+#         # save api key
+#         st.session_state.openai_api_key = openai_api_key_st
+
 #### load builder agent and its tool spec (the agent_builder)
 builder_agent, agent_builder = load_meta_agent_and_tools()
 

README.md

@@ -17,6 +17,15 @@ Clone this project, go into the `rags` project folder.
 pip install -r requirements.txt
 ```
 
+By default, we use OpenAI for both the builder agent as well as the generated RAG agent.
+Please `.streamlit/secrets.toml` in the home folder.
+
+Then put the following:
+```
+openai_key = "<openai_key>"
+```
+
+
 Then run the app from the "home page" file.
 
 ```
@@ -64,6 +73,24 @@ This is a standard chatbot interface where you can query the RAG agent and it wi
 It will be able to pick the right RAG tools (either top-k vector search or optionally summarization) in order to fulfill the query.
 
 
+## Supported LLMs and Embeddings
+
+### Builder Agent
+
+By default the builder agent uses OpenAI. This is defined in the `builder_config.py` file.
+
+You can customize this to whatever LLM you want (an example is provided for Anthropic).
+
+Note that GPT-4 variants will give the most reliable results in terms of actually constructing an agent (we couldn't get Claude to work).
+
+### Generated RAG Agent
+
+You can set the configuration either through natural language or manually for both the embedding model and LLM.
+
+- **LLM**: Currently only OpenAI LLMs are supported (conceptually we can support other LLMs but just need to figure out how to map names to LLM configurations)
+- **Embeddings**: Supports text-embedding-ada-002 by default, but also supports Hugging Face models. To use a hugging face model simply prepend with local, e.g. local:BAAI/bge-small-en.
+
+
 ## Issues / Contributions
 
 Running into issues? Please file a Github issue or join our [Discord](https://discord.gg/dGcwcsnxhU).

agent_utils.py

@@ -1,9 +1,10 @@
 from llama_index.llms import OpenAI, ChatMessage
+from llama_index.llms.base import LLM
 from llama_index.llms.utils import resolve_llm
 from pydantic import BaseModel, Field
 import os
 from llama_index.tools.query_engine import QueryEngineTool
-from llama_index.agent import OpenAIAgent
+from llama_index.agent import OpenAIAgent, ReActAgent
 from llama_index import (
     VectorStoreIndex,
     SummaryIndex,
@@ -15,8 +16,14 @@
 from llama_index import SimpleDirectoryReader
 from llama_index.embeddings.utils import resolve_embed_model
 from llama_index.tools import QueryEngineTool, ToolMetadata, FunctionTool
-from typing import Dict, Tuple
+from llama_index.agent.types import BaseAgent
+from llama_index.agent.react.formatter import ReActChatFormatter
+from llama_index.llms.openai_utils import is_function_calling_model
+from builder_config import BUILDER_LLM
+from typing import Dict, Tuple, Any
 import streamlit as st
+from pathlib import Path
+import json
 
 
 ####################
@@ -50,6 +57,33 @@
 GEN_SYS_PROMPT_TMPL = ChatPromptTemplate(gen_sys_prompt_messages)
 
 
+def load_agent(
+    tools: List, 
+    llm: LLM, 
+    system_prompt: str,
+    **kwargs: Any
+) -> BaseAgent:
+    """Load agent."""
+    if isinstance(llm, OpenAI) and is_function_calling_model(llm.model):
+        # get OpenAI Agent
+        agent = OpenAIAgent.from_tools(
+            tools=tools,
+            llm=llm,
+            system_prompt=system_prompt,
+            **kwargs
+        )
+    else:
+        agent = ReActAgent.from_tools(
+            tools=tools,
+            llm=llm,
+            react_chat_formatter=ReActChatFormatter(
+                system_header=system_prompt,
+            ),
+            **kwargs
+        )
+    return agent
+
+
 class RAGParams(BaseModel):
     """RAG parameters.
 
@@ -252,12 +286,10 @@ def create_agent(self) -> None:
         if self._cache.system_prompt is None:
             return "System prompt not set yet. Please set system prompt first."
 
-        agent = OpenAIAgent.from_tools(
-            tools=all_tools,
-            system_prompt=self._cache.system_prompt,
-            llm=llm,
-            verbose=True
+        agent = load_agent(
+            all_tools, llm=llm, system_prompt=self._cache.system_prompt, verbose=True
         )
+
         self._cache.agent = agent
         return "Agent created successfully."
 
@@ -284,10 +316,14 @@ def create_agent(self) -> None:
 """
 
 
+### DEFINE Agent ####
+# NOTE: here we define a function that is dependent on the LLM,
+# please make sure to update the LLM above if you change the function below
+
+
 # define agent
 @st.cache_resource
 def load_meta_agent_and_tools() -> Tuple[OpenAIAgent, RAGAgentBuilder]:
-    prefix_msgs = [ChatMessage(role="system", content=RAG_BUILDER_SYS_STR)]
 
     # think of this as tools for the agent to use
     agent_builder = RAGAgentBuilder()
@@ -302,11 +338,9 @@ def load_meta_agent_and_tools() -> Tuple[OpenAIAgent, RAGAgentBuilder]:
     ]
     fn_tools = [FunctionTool.from_defaults(fn=fn) for fn in fns]
 
-    builder_agent = OpenAIAgent.from_tools(
-        tools=fn_tools,
-        llm=OpenAI(llm="gpt-4-1106-preview"),
-        prefix_messages=prefix_msgs,
-        verbose=True,
+    builder_agent = load_agent(
+        fn_tools, llm=BUILDER_LLM, system_prompt=RAG_BUILDER_SYS_STR, verbose=True
     )
+
     return builder_agent, agent_builder
 

builder_config.py

@@ -0,0 +1,20 @@
+"""Configuration."""
+import streamlit as st
+import os
+
+### DEFINE BUILDER_LLM #####
+## Uncomment the LLM you want to use to construct the meta agent
+
+## OpenAI
+from llama_index.llms import OpenAI
+# set OpenAI Key - use Streamlit secrets
+os.environ["OPENAI_API_KEY"] = st.secrets.openai_key
+# load LLM
+BUILDER_LLM = OpenAI(model="gpt-4-1106-preview")
+
+## Anthropic (make sure you `pip install anthropic`)
+## NOTE: Hallucinates
+# from llama_index.llms import Anthropic
+# # set Anthropic key
+# os.environ["ANTHROPIC_API_KEY"] = st.secrets.anthropic_key
+# BUILDER_LLM = Anthropic()

pages/2_⚙️_RAG_Config.py

@@ -3,7 +3,6 @@
 import openai
 from streamlit_pills import pills
 from typing import cast
-from llama_index.agent.openai_agent import OpenAIAgent
 
 from agent_utils import (
     RAGParams,
@@ -22,7 +21,6 @@
 st.info(
     "This is generated by the builder in the above section.", icon="ℹ️"
 )
-openai.api_key = st.secrets.openai_key
 
 if "agent_builder" in st.session_state.keys():
     agent_builder = cast(RAGAgentBuilder, st.session_state.agent_builder)

pages/3_🤖_Generated_RAG_Agent.py

@@ -1,11 +1,6 @@
 """Streamlit page showing builder config."""
-import os
 import streamlit as st
-os.environ["OPENAI_API_KEY"] = st.secrets.openai_key
-import openai
-from streamlit_pills import pills
 from typing import cast
-from llama_index.agent.openai_agent import OpenAIAgent
 from agent_utils import (
     RAGAgentBuilder,
 )

requirements.txt

@@ -2,4 +2,5 @@ streamlit
 streamlit_pills
 llama-index==0.9.4
 llama-hub==0.0.44
+# NOTE: this is due to a trivial dependency in the web tool, will refactor
 langchain==0.0.305