From 04cdaa112fed6b68913bfeba090bad9d9e8813a7 Mon Sep 17 00:00:00 2001 From: Jerry Liu Date: Sat, 2 Dec 2023 00:04:26 -0800 Subject: [PATCH] cr --- "1_\360\237\217\240_Home.py" | 52 ++++++++----------- core/agent_builder.py | 19 ++----- core/param_cache.py | 25 --------- core/utils.py | 14 ++--- .../2_\342\232\231\357\270\217_RAG_Config.py" | 8 +-- ...3_\360\237\244\226_Generated_RAG_Agent.py" | 4 -- 6 files changed, 32 insertions(+), 90 deletions(-) diff --git "a/1_\360\237\217\240_Home.py" "b/1_\360\237\217\240_Home.py" index 0b2bbbd..e32e921 100644 --- "a/1_\360\237\217\240_Home.py" +++ "b/1_\360\237\217\240_Home.py" @@ -1,16 +1,7 @@ import streamlit as st from streamlit_pills import pills -from typing import cast -from core.agent_builder import ( - load_meta_agent_and_tools, - load_agent_ids_from_directory, - AgentCacheRegistry, -) from st_utils import add_sidebar, get_current_state -from core.constants import ( - AGENT_CACHE_DIR, -) #################### @@ -77,24 +68,25 @@ def add_to_message_history(role: str, content: str) -> None: with st.chat_message("user"): st.write(prompt) -# If last message is not from assistant, generate a new response -if st.session_state.messages[-1]["role"] != "assistant": - with st.chat_message("assistant"): - with st.spinner("Thinking..."): - response = current_state.builder_agent.chat(prompt) - st.write(str(response)) - add_to_message_history("assistant", str(response)) - - # check agent_ids again, if it doesn't match, add to directory and refresh - agent_ids = current_state.agent_registry.get_agent_ids() - # check diff between agent_ids and cur agent ids - diff_ids = list(set(agent_ids) - set(st.session_state.cur_agent_ids)) - if len(diff_ids) > 0: - # clear streamlit cache, to allow you to generate a new agent - st.cache_resource.clear() - - # trigger refresh - st.rerun() - -else: - pass + # If last message is not from assistant, generate a new response + if st.session_state.messages[-1]["role"] != "assistant": + with st.chat_message("assistant"): + with st.spinner("Thinking..."): + response = current_state.builder_agent.chat(prompt) + st.write(str(response)) + add_to_message_history("assistant", str(response)) + + # check agent_ids again + # if it doesn't match, add to directory and refresh + agent_ids = current_state.agent_registry.get_agent_ids() + # check diff between agent_ids and cur agent ids + diff_ids = list(set(agent_ids) - set(st.session_state.cur_agent_ids)) + if len(diff_ids) > 0: + # clear streamlit cache, to allow you to generate a new agent + st.cache_resource.clear() + + # trigger refresh + st.rerun() + + else: + pass diff --git a/core/agent_builder.py b/core/agent_builder.py index 553853b..95f02dc 100644 --- a/core/agent_builder.py +++ b/core/agent_builder.py @@ -1,19 +1,10 @@ """Agent builder.""" -from llama_index.llms import OpenAI, ChatMessage, Anthropic, Replicate -from llama_index.llms.base import LLM -from llama_index.llms.utils import resolve_llm -from pydantic import BaseModel, Field +from llama_index.llms import ChatMessage from llama_index.prompts import ChatPromptTemplate from typing import List, cast, Optional -from llama_index import SimpleDirectoryReader -from llama_index.embeddings.utils import resolve_embed_model -from llama_index.tools import QueryEngineTool, ToolMetadata, FunctionTool +from llama_index.tools import FunctionTool from llama_index.agent.types import BaseAgent -from llama_index.chat_engine.types import BaseChatEngine -from llama_index.agent.react.formatter import ReActChatFormatter -from llama_index.llms.openai_utils import is_function_calling_model -from llama_index.chat_engine import CondensePlusContextChatEngine from core.builder_config import BUILDER_LLM from typing import Dict, Tuple, Any, Callable, Union import streamlit as st @@ -23,8 +14,6 @@ from core.constants import AGENT_CACHE_DIR import shutil -from llama_index.callbacks import CallbackManager -from callback_manager import StreamlitFunctionsCallbackHandler from core.param_cache import ParamCache, RAGParams from core.utils import ( load_data, @@ -152,7 +141,9 @@ def __init__( ) -> None: """Init params.""" self._cache = cache or ParamCache() - self._agent_registry = agent_registry or AgentCacheRegistry(str(AGENT_CACHE_DIR)) + self._agent_registry = agent_registry or AgentCacheRegistry( + str(AGENT_CACHE_DIR) + ) @property def cache(self) -> ParamCache: diff --git a/core/param_cache.py b/core/param_cache.py index 4a0e748..a7cd6e8 100644 --- a/core/param_cache.py +++ b/core/param_cache.py @@ -1,43 +1,18 @@ """Param cache.""" -from llama_index.llms import OpenAI, ChatMessage, Anthropic, Replicate -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.agent import OpenAIAgent, ReActAgent -from llama_index.agent.react.prompts import REACT_CHAT_SYSTEM_HEADER from llama_index import ( VectorStoreIndex, - SummaryIndex, - ServiceContext, StorageContext, - Document, load_index_from_storage, ) -from llama_index.prompts import ChatPromptTemplate from typing import List, cast, Optional -from llama_index import SimpleDirectoryReader -from llama_index.embeddings.utils import resolve_embed_model -from llama_index.tools import QueryEngineTool, ToolMetadata, FunctionTool -from llama_index.agent.types import BaseAgent from llama_index.chat_engine.types import BaseChatEngine -from llama_index.agent.react.formatter import ReActChatFormatter -from llama_index.llms.openai_utils import is_function_calling_model -from llama_index.chat_engine import CondensePlusContextChatEngine -from core.builder_config import BUILDER_LLM -from typing import Dict, Tuple, Any, Callable -import streamlit as st from pathlib import Path import json import uuid -from core.constants import AGENT_CACHE_DIR -import shutil from core.utils import load_data, get_tool_objects, construct_agent, RAGParams -from llama_index.callbacks import CallbackManager -from callback_manager import StreamlitFunctionsCallbackHandler - class ParamCache(BaseModel): """Cache for RAG agent builder. diff --git a/core/utils.py b/core/utils.py index ecb03ee..639e86e 100644 --- a/core/utils.py +++ b/core/utils.py @@ -1,6 +1,6 @@ """Utils.""" -from llama_index.llms import OpenAI, ChatMessage, Anthropic, Replicate +from llama_index.llms import OpenAI, Anthropic, Replicate from llama_index.llms.base import LLM from llama_index.llms.utils import resolve_llm from pydantic import BaseModel, Field @@ -11,28 +11,20 @@ VectorStoreIndex, SummaryIndex, ServiceContext, - StorageContext, Document, - load_index_from_storage, ) -from llama_index.prompts import ChatPromptTemplate from typing import List, cast, Optional from llama_index import SimpleDirectoryReader from llama_index.embeddings.utils import resolve_embed_model -from llama_index.tools import QueryEngineTool, ToolMetadata, FunctionTool +from llama_index.tools import QueryEngineTool, ToolMetadata from llama_index.agent.types import BaseAgent from llama_index.chat_engine.types import BaseChatEngine from llama_index.agent.react.formatter import ReActChatFormatter from llama_index.llms.openai_utils import is_function_calling_model from llama_index.chat_engine import CondensePlusContextChatEngine from core.builder_config import BUILDER_LLM -from typing import Dict, Tuple, Any, Callable +from typing import Dict, Tuple, Any import streamlit as st -from pathlib import Path -import json -import uuid -from core.constants import AGENT_CACHE_DIR -import shutil from llama_index.callbacks import CallbackManager from callback_manager import StreamlitFunctionsCallbackHandler diff --git "a/pages/2_\342\232\231\357\270\217_RAG_Config.py" "b/pages/2_\342\232\231\357\270\217_RAG_Config.py" index d68ed50..3f36daa 100644 --- "a/pages/2_\342\232\231\357\270\217_RAG_Config.py" +++ "b/pages/2_\342\232\231\357\270\217_RAG_Config.py" @@ -1,19 +1,15 @@ """Streamlit page showing builder config.""" import streamlit as st -from typing import cast, Optional from core.param_cache import ( RAGParams, - ParamCache, ) from core.agent_builder import ( RAGAgentBuilder, - remove_agent_from_directory, AgentCacheRegistry, ) -from st_utils import update_selected_agent_with_id, get_current_state -from core.constants import AGENT_CACHE_DIR -from st_utils import add_sidebar +from st_utils import update_selected_agent_with_id, get_current_state, add_sidebar +from typing import cast #################### diff --git "a/pages/3_\360\237\244\226_Generated_RAG_Agent.py" "b/pages/3_\360\237\244\226_Generated_RAG_Agent.py" index 703ffea..b0b83bb 100644 --- "a/pages/3_\360\237\244\226_Generated_RAG_Agent.py" +++ "b/pages/3_\360\237\244\226_Generated_RAG_Agent.py" @@ -1,9 +1,5 @@ """Streamlit page showing builder config.""" import streamlit as st -from typing import cast, Optional -from core.agent_builder import RAGAgentBuilder, AgentCacheRegistry -from core.param_cache import ParamCache -from core.constants import AGENT_CACHE_DIR from st_utils import add_sidebar, get_current_state