uploaded all small directories

core/__init__.py

@@ -0,0 +1 @@
+"""Init file."""

core/agent_builder/base.py

@@ -0,0 +1,250 @@
+"""Agent builder."""
+
+from llama_index.llms import ChatMessage
+from llama_index.prompts import ChatPromptTemplate
+from typing import List, cast, Optional
+from core.builder_config import BUILDER_LLM
+from typing import Dict, Any
+import uuid
+from core.constants import AGENT_CACHE_DIR
+from abc import ABC, abstractmethod
+
+from core.param_cache import ParamCache, RAGParams
+from core.utils import (
+    load_data,
+    get_tool_objects,
+    construct_agent,
+)
+from core.agent_builder.registry import AgentCacheRegistry
+
+
+# System prompt tool
+GEN_SYS_PROMPT_STR = """\
+Task information is given below. 
+
+Given the task, please generate a system prompt for an OpenAI-powered bot \
+to solve this task: 
+{task} \
+
+Make sure the system prompt obeys the following requirements:
+- Tells the bot to ALWAYS use tools given to solve the task. \
+NEVER give an answer without using a tool.
+- Does not reference a specific data source. \
+The data source is implicit in any queries to the bot, \
+and telling the bot to analyze a specific data source might confuse it given a \
+user query.
+
+"""
+
+gen_sys_prompt_messages = [
+    ChatMessage(
+        role="system",
+        content="You are helping to build a system prompt for another bot.",
+    ),
+    ChatMessage(role="user", content=GEN_SYS_PROMPT_STR),
+]
+
+GEN_SYS_PROMPT_TMPL = ChatPromptTemplate(gen_sys_prompt_messages)
+
+
+class BaseRAGAgentBuilder(ABC):
+    """Base RAG Agent builder class."""
+
+    @property
+    @abstractmethod
+    def cache(self) -> ParamCache:
+        """Cache."""
+
+    @property
+    @abstractmethod
+    def agent_registry(self) -> AgentCacheRegistry:
+        """Agent registry."""
+
+
+class RAGAgentBuilder(BaseRAGAgentBuilder):
+    """RAG Agent builder.
+
+    Contains a set of functions to construct a RAG agent, including:
+    - setting system prompts
+    - loading data
+    - adding web search
+    - setting parameters (e.g. top-k)
+
+    Must pass in a cache. This cache will be modified as the agent is built.
+
+    """
+
+    def __init__(
+        self,
+        cache: Optional[ParamCache] = None,
+        agent_registry: Optional[AgentCacheRegistry] = None,
+    ) -> None:
+        """Init params."""
+        self._cache = cache or ParamCache()
+        self._agent_registry = agent_registry or AgentCacheRegistry(
+            str(AGENT_CACHE_DIR)
+        )
+
+    @property
+    def cache(self) -> ParamCache:
+        """Cache."""
+        return self._cache
+
+    @property
+    def agent_registry(self) -> AgentCacheRegistry:
+        """Agent registry."""
+        return self._agent_registry
+
+    def create_system_prompt(self, task: str) -> str:
+        """Create system prompt for another agent given an input task."""
+        llm = BUILDER_LLM
+        fmt_messages = GEN_SYS_PROMPT_TMPL.format_messages(task=task)
+        response = llm.chat(fmt_messages)
+        self._cache.system_prompt = response.message.content
+
+        return f"System prompt created: {response.message.content}"
+
+    def load_data(
+        self,
+        file_names: Optional[List[str]] = None,
+        directory: Optional[str] = None,
+        urls: Optional[List[str]] = None,
+    ) -> str:
+        """Load data for a given task.
+
+        Only ONE of file_names or directory or urls should be specified.
+
+        Args:
+            file_names (Optional[List[str]]): List of file names to load.
+                Defaults to None.
+            directory (Optional[str]): Directory to load files from.
+            urls (Optional[List[str]]): List of urls to load.
+                Defaults to None.
+
+        """
+        file_names = file_names or []
+        urls = urls or []
+        directory = directory or ""
+        docs = load_data(file_names=file_names, directory=directory, urls=urls)
+        self._cache.docs = docs
+        self._cache.file_names = file_names
+        self._cache.urls = urls
+        self._cache.directory = directory
+        return "Data loaded successfully."
+
+    def add_web_tool(self) -> str:
+        """Add a web tool to enable agent to solve a task."""
+        # TODO: make this not hardcoded to a web tool
+        # Set up Metaphor tool
+        if "web_search" in self._cache.tools:
+            return "Web tool already added."
+        else:
+            self._cache.tools.append("web_search")
+        return "Web tool added successfully."
+
+    def get_rag_params(self) -> Dict:
+        """Get parameters used to configure the RAG pipeline.
+
+        Should be called before `set_rag_params` so that the agent is aware of the
+        schema.
+
+        """
+        rag_params = self._cache.rag_params
+        return rag_params.dict()
+
+    def set_rag_params(self, **rag_params: Dict) -> str:
+        """Set RAG parameters.
+
+        These parameters will then be used to actually initialize the agent.
+        Should call `get_rag_params` first to get the schema of the input dictionary.
+
+        Args:
+            **rag_params (Dict): dictionary of RAG parameters.
+
+        """
+        new_dict = self._cache.rag_params.dict()
+        new_dict.update(rag_params)
+        rag_params_obj = RAGParams(**new_dict)
+        self._cache.rag_params = rag_params_obj
+        return "RAG parameters set successfully."
+
+    def create_agent(self, agent_id: Optional[str] = None) -> str:
+        """Create an agent.
+
+        There are no parameters for this function because all the
+        functions should have already been called to set up the agent.
+
+        """
+        if self._cache.system_prompt is None:
+            raise ValueError("Must set system prompt before creating agent.")
+
+        # construct additional tools
+        additional_tools = get_tool_objects(self.cache.tools)
+        agent, extra_info = construct_agent(
+            cast(str, self._cache.system_prompt),
+            cast(RAGParams, self._cache.rag_params),
+            self._cache.docs,
+            additional_tools=additional_tools,
+        )
+
+        # if agent_id not specified, randomly generate one
+        agent_id = agent_id or self._cache.agent_id or f"Agent_{str(uuid.uuid4())}"
+        self._cache.vector_index = extra_info["vector_index"]
+        self._cache.agent_id = agent_id
+        self._cache.agent = agent
+
+        # save the cache to disk
+        self._agent_registry.add_new_agent_cache(agent_id, self._cache)
+        return "Agent created successfully."
+
+    def update_agent(
+        self,
+        agent_id: str,
+        system_prompt: Optional[str] = None,
+        include_summarization: Optional[bool] = None,
+        top_k: Optional[int] = None,
+        chunk_size: Optional[int] = None,
+        embed_model: Optional[str] = None,
+        llm: Optional[str] = None,
+        additional_tools: Optional[List] = None,
+    ) -> None:
+        """Update agent.
+
+        Delete old agent by ID and create a new one.
+        Optionally update the system prompt and RAG parameters.
+
+        NOTE: Currently is manually called, not meant for agent use.
+
+        """
+        self._agent_registry.delete_agent_cache(self.cache.agent_id)
+
+        # set agent id
+        self.cache.agent_id = agent_id
+
+        # set system prompt
+        if system_prompt is not None:
+            self.cache.system_prompt = system_prompt
+        # get agent_builder
+        # We call set_rag_params and create_agent, which will
+        # update the cache
+        # TODO: decouple functions from tool functions exposed to the agent
+        rag_params_dict: Dict[str, Any] = {}
+        if include_summarization is not None:
+            rag_params_dict["include_summarization"] = include_summarization
+        if top_k is not None:
+            rag_params_dict["top_k"] = top_k
+        if chunk_size is not None:
+            rag_params_dict["chunk_size"] = chunk_size
+        if embed_model is not None:
+            rag_params_dict["embed_model"] = embed_model
+        if llm is not None:
+            rag_params_dict["llm"] = llm
+
+        self.set_rag_params(**rag_params_dict)
+
+        # update tools
+        if additional_tools is not None:
+            self.cache.tools = additional_tools
+
+        # this will update the agent in the cache
+        self.create_agent()

core/agent_builder/loader.py

@@ -0,0 +1,115 @@
+"""Loader agent."""
+
+from typing import List, cast, Optional
+from llama_index.tools import FunctionTool
+from llama_index.agent.types import BaseAgent
+from core.builder_config import BUILDER_LLM
+from typing import Tuple, Callable
+import streamlit as st
+
+from core.param_cache import ParamCache
+from core.utils import (
+    load_meta_agent,
+)
+from core.agent_builder.registry import AgentCacheRegistry
+from core.agent_builder.base import RAGAgentBuilder, BaseRAGAgentBuilder
+from core.agent_builder.multimodal import MultimodalRAGAgentBuilder
+
+####################
+#### META Agent ####
+####################
+
+RAG_BUILDER_SYS_STR = """\
+You are helping to construct an agent given a user-specified task. 
+You should generally use the tools in this rough order to build the agent.
+
+1) Create system prompt tool: to create the system prompt for the agent.
+2) Load in user-specified data (based on file paths they specify).
+3) Decide whether or not to add additional tools.
+4) Set parameters for the RAG pipeline.
+5) Build the agent
+
+This will be a back and forth conversation with the user. You should
+continue asking users if there's anything else they want to do until
+they say they're done. To help guide them on the process, 
+you can give suggestions on parameters they can set based on the tools they
+have available (e.g. "Do you want to set the number of documents to retrieve?")
+
+"""
+
+
+### 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
+
+
+def _get_builder_agent_tools(agent_builder: RAGAgentBuilder) -> List[FunctionTool]:
+    """Get list of builder agent tools to pass to the builder agent."""
+    # see if metaphor api key is set, otherwise don't add web tool
+    # TODO: refactor this later
+
+    if "metaphor_key" in st.secrets:
+        fns: List[Callable] = [
+            agent_builder.create_system_prompt,
+            agent_builder.load_data,
+            agent_builder.add_web_tool,
+            agent_builder.get_rag_params,
+            agent_builder.set_rag_params,
+            agent_builder.create_agent,
+        ]
+    else:
+        fns = [
+            agent_builder.create_system_prompt,
+            agent_builder.load_data,
+            agent_builder.get_rag_params,
+            agent_builder.set_rag_params,
+            agent_builder.create_agent,
+        ]
+
+    fn_tools: List[FunctionTool] = [FunctionTool.from_defaults(fn=fn) for fn in fns]
+    return fn_tools
+
+
+def _get_mm_builder_agent_tools(
+    agent_builder: MultimodalRAGAgentBuilder,
+) -> List[FunctionTool]:
+    """Get list of builder agent tools to pass to the builder agent."""
+    fns: List[Callable] = [
+        agent_builder.create_system_prompt,
+        agent_builder.load_data,
+        agent_builder.get_rag_params,
+        agent_builder.set_rag_params,
+        agent_builder.create_agent,
+    ]
+
+    fn_tools: List[FunctionTool] = [FunctionTool.from_defaults(fn=fn) for fn in fns]
+    return fn_tools
+
+
+# define agent
+def load_meta_agent_and_tools(
+    cache: Optional[ParamCache] = None,
+    agent_registry: Optional[AgentCacheRegistry] = None,
+    is_multimodal: bool = False,
+) -> Tuple[BaseAgent, BaseRAGAgentBuilder]:
+    """Load meta agent and tools."""
+
+    if is_multimodal:
+        agent_builder: BaseRAGAgentBuilder = MultimodalRAGAgentBuilder(
+            cache, agent_registry=agent_registry
+        )
+        fn_tools = _get_mm_builder_agent_tools(
+            cast(MultimodalRAGAgentBuilder, agent_builder)
+        )
+        builder_agent = load_meta_agent(
+            fn_tools, llm=BUILDER_LLM, system_prompt=RAG_BUILDER_SYS_STR, verbose=True
+        )
+    else:
+        # think of this as tools for the agent to use
+        agent_builder = RAGAgentBuilder(cache, agent_registry=agent_registry)
+        fn_tools = _get_builder_agent_tools(agent_builder)
+        builder_agent = load_meta_agent(
+            fn_tools, llm=BUILDER_LLM, system_prompt=RAG_BUILDER_SYS_STR, verbose=True
+        )
+
+    return builder_agent, agent_builder

core/agent_builder/multimodal.py

@@ -0,0 +1,256 @@
+"""Multimodal agent builder."""
+
+from llama_index.llms import ChatMessage
+from typing import List, cast, Optional
+from core.builder_config import BUILDER_LLM
+from typing import Dict, Any
+import uuid
+from core.constants import AGENT_CACHE_DIR
+
+from core.param_cache import ParamCache, RAGParams
+from core.utils import (
+    load_data,
+    construct_mm_agent,
+)
+from core.agent_builder.registry import AgentCacheRegistry
+from core.agent_builder.base import GEN_SYS_PROMPT_TMPL, BaseRAGAgentBuilder
+
+from llama_index.chat_engine.types import BaseChatEngine
+
+from llama_index.callbacks import trace_method
+from llama_index.query_engine.multi_modal import SimpleMultiModalQueryEngine
+from llama_index.chat_engine.types import (
+    AGENT_CHAT_RESPONSE_TYPE,
+    StreamingAgentChatResponse,
+    AgentChatResponse,
+)
+from llama_index.llms.base import ChatResponse
+from typing import Generator
+
+
+class MultimodalChatEngine(BaseChatEngine):
+    """Multimodal chat engine.
+
+    This chat engine is a light wrapper around a query engine.
+    Offers no real 'chat' functionality, is a beta feature.
+
+    """
+
+    def __init__(self, mm_query_engine: SimpleMultiModalQueryEngine) -> None:
+        """Init params."""
+        self._mm_query_engine = mm_query_engine
+
+    def reset(self) -> None:
+        """Reset conversation state."""
+        pass
+
+    @trace_method("chat")
+    def chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> AGENT_CHAT_RESPONSE_TYPE:
+        """Main chat interface."""
+        # just return the top-k results
+        response = self._mm_query_engine.query(message)
+        return AgentChatResponse(response=str(response))
+
+    @trace_method("chat")
+    def stream_chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> StreamingAgentChatResponse:
+        """Stream chat interface."""
+        response = self._mm_query_engine.query(message)
+
+        def _chat_stream(response: str) -> Generator[ChatResponse, None, None]:
+            yield ChatResponse(message=ChatMessage(role="assistant", content=response))
+
+        chat_stream = _chat_stream(str(response))
+        return StreamingAgentChatResponse(chat_stream=chat_stream)
+
+    @trace_method("chat")
+    async def achat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> AGENT_CHAT_RESPONSE_TYPE:
+        """Async version of main chat interface."""
+        response = await self._mm_query_engine.aquery(message)
+        return AgentChatResponse(response=str(response))
+
+    @trace_method("chat")
+    async def astream_chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> StreamingAgentChatResponse:
+        """Async version of main chat interface."""
+        return self.stream_chat(message, chat_history)
+
+
+class MultimodalRAGAgentBuilder(BaseRAGAgentBuilder):
+    """Multimodal RAG Agent builder.
+
+    Contains a set of functions to construct a RAG agent, including:
+    - setting system prompts
+    - loading data
+    - adding web search
+    - setting parameters (e.g. top-k)
+
+    Must pass in a cache. This cache will be modified as the agent is built.
+
+    """
+
+    def __init__(
+        self,
+        cache: Optional[ParamCache] = None,
+        agent_registry: Optional[AgentCacheRegistry] = None,
+    ) -> None:
+        """Init params."""
+        self._cache = cache or ParamCache()
+        self._agent_registry = agent_registry or AgentCacheRegistry(
+            str(AGENT_CACHE_DIR)
+        )
+
+    @property
+    def cache(self) -> ParamCache:
+        """Cache."""
+        return self._cache
+
+    @property
+    def agent_registry(self) -> AgentCacheRegistry:
+        """Agent registry."""
+        return self._agent_registry
+
+    def create_system_prompt(self, task: str) -> str:
+        """Create system prompt for another agent given an input task."""
+        llm = BUILDER_LLM
+        fmt_messages = GEN_SYS_PROMPT_TMPL.format_messages(task=task)
+        response = llm.chat(fmt_messages)
+        self._cache.system_prompt = response.message.content
+
+        return f"System prompt created: {response.message.content}"
+
+    def load_data(
+        self,
+        file_names: Optional[List[str]] = None,
+        directory: Optional[str] = None,
+    ) -> str:
+        """Load data for a given task.
+
+        Only ONE of file_names or directory should be specified.
+        **NOTE**: urls not supported in multi-modal setting.
+
+        Args:
+            file_names (Optional[List[str]]): List of file names to load.
+                Defaults to None.
+            directory (Optional[str]): Directory to load files from.
+
+        """
+        file_names = file_names or []
+        directory = directory or ""
+        docs = load_data(file_names=file_names, directory=directory)
+        self._cache.docs = docs
+        self._cache.file_names = file_names
+        self._cache.directory = directory
+        return "Data loaded successfully."
+
+    def get_rag_params(self) -> Dict:
+        """Get parameters used to configure the RAG pipeline.
+
+        Should be called before `set_rag_params` so that the agent is aware of the
+        schema.
+
+        """
+        rag_params = self._cache.rag_params
+        return rag_params.dict()
+
+    def set_rag_params(self, **rag_params: Dict) -> str:
+        """Set RAG parameters.
+
+        These parameters will then be used to actually initialize the agent.
+        Should call `get_rag_params` first to get the schema of the input dictionary.
+
+        Args:
+            **rag_params (Dict): dictionary of RAG parameters.
+
+        """
+        new_dict = self._cache.rag_params.dict()
+        new_dict.update(rag_params)
+        rag_params_obj = RAGParams(**new_dict)
+        self._cache.rag_params = rag_params_obj
+        return "RAG parameters set successfully."
+
+    def create_agent(self, agent_id: Optional[str] = None) -> str:
+        """Create an agent.
+
+        There are no parameters for this function because all the
+        functions should have already been called to set up the agent.
+
+        """
+        if self._cache.system_prompt is None:
+            raise ValueError("Must set system prompt before creating agent.")
+
+        # construct additional tools
+        agent, extra_info = construct_mm_agent(
+            cast(str, self._cache.system_prompt),
+            cast(RAGParams, self._cache.rag_params),
+            self._cache.docs,
+        )
+
+        # if agent_id not specified, randomly generate one
+        agent_id = agent_id or self._cache.agent_id or f"Agent_{str(uuid.uuid4())}"
+        self._cache.builder_type = "multimodal"
+        self._cache.vector_index = extra_info["vector_index"]
+        self._cache.agent_id = agent_id
+        self._cache.agent = agent
+
+        # save the cache to disk
+        self._agent_registry.add_new_agent_cache(agent_id, self._cache)
+        return "Agent created successfully."
+
+    def update_agent(
+        self,
+        agent_id: str,
+        system_prompt: Optional[str] = None,
+        include_summarization: Optional[bool] = None,
+        top_k: Optional[int] = None,
+        chunk_size: Optional[int] = None,
+        embed_model: Optional[str] = None,
+        llm: Optional[str] = None,
+        additional_tools: Optional[List] = None,
+    ) -> None:
+        """Update agent.
+
+        Delete old agent by ID and create a new one.
+        Optionally update the system prompt and RAG parameters.
+
+        NOTE: Currently is manually called, not meant for agent use.
+
+        """
+        self._agent_registry.delete_agent_cache(self.cache.agent_id)
+
+        # set agent id
+        self.cache.agent_id = agent_id
+
+        # set system prompt
+        if system_prompt is not None:
+            self.cache.system_prompt = system_prompt
+        # get agent_builder
+        # We call set_rag_params and create_agent, which will
+        # update the cache
+        # TODO: decouple functions from tool functions exposed to the agent
+        rag_params_dict: Dict[str, Any] = {}
+        if include_summarization is not None:
+            rag_params_dict["include_summarization"] = include_summarization
+        if top_k is not None:
+            rag_params_dict["top_k"] = top_k
+        if chunk_size is not None:
+            rag_params_dict["chunk_size"] = chunk_size
+        if embed_model is not None:
+            rag_params_dict["embed_model"] = embed_model
+        if llm is not None:
+            rag_params_dict["llm"] = llm
+
+        self.set_rag_params(**rag_params_dict)
+
+        # update tools
+        if additional_tools is not None:
+            self.cache.tools = additional_tools
+
+        # this will update the agent in the cache
+        self.create_agent()

core/agent_builder/registry.py

@@ -0,0 +1,78 @@
+"""Agent builder registry."""
+
+from typing import List
+from typing import Union
+from pathlib import Path
+import json
+import shutil
+
+from core.param_cache import ParamCache
+
+
+class AgentCacheRegistry:
+    """Registry for agent caches, in disk.
+
+    Can register new agent caches, load agent caches, delete agent caches, etc.
+
+    """
+
+    def __init__(self, dir: Union[str, Path]) -> None:
+        """Init params."""
+        self._dir = dir
+
+    def _add_agent_id_to_directory(self, agent_id: str) -> None:
+        """Save agent id to directory."""
+        full_path = Path(self._dir) / "agent_ids.json"
+        if not full_path.exists():
+            with open(full_path, "w") as f:
+                json.dump({"agent_ids": [agent_id]}, f)
+        else:
+            with open(full_path, "r") as f:
+                agent_ids = json.load(f)["agent_ids"]
+            if agent_id in agent_ids:
+                raise ValueError(f"Agent id {agent_id} already exists.")
+            agent_ids_set = set(agent_ids)
+            agent_ids_set.add(agent_id)
+            with open(full_path, "w") as f:
+                json.dump({"agent_ids": list(agent_ids_set)}, f)
+
+    def add_new_agent_cache(self, agent_id: str, cache: ParamCache) -> None:
+        """Register agent."""
+        # save the cache to disk
+        agent_cache_path = f"{self._dir}/{agent_id}"
+        cache.save_to_disk(agent_cache_path)
+        # save to agent ids
+        self._add_agent_id_to_directory(agent_id)
+
+    def get_agent_ids(self) -> List[str]:
+        """Get agent ids."""
+        full_path = Path(self._dir) / "agent_ids.json"
+        if not full_path.exists():
+            return []
+        with open(full_path, "r") as f:
+            agent_ids = json.load(f)["agent_ids"]
+
+        return agent_ids
+
+    def get_agent_cache(self, agent_id: str) -> ParamCache:
+        """Get agent cache."""
+        full_path = Path(self._dir) / f"{agent_id}"
+        if not full_path.exists():
+            raise ValueError(f"Cache for agent {agent_id} does not exist.")
+        cache = ParamCache.load_from_disk(str(full_path))
+        return cache
+
+    def delete_agent_cache(self, agent_id: str) -> None:
+        """Delete agent cache."""
+        # modify / resave agent_ids
+        agent_ids = self.get_agent_ids()
+        new_agent_ids = [id for id in agent_ids if id != agent_id]
+        full_path = Path(self._dir) / "agent_ids.json"
+        with open(full_path, "w") as f:
+            json.dump({"agent_ids": new_agent_ids}, f)
+
+        # remove agent cache
+        full_path = Path(self._dir) / f"{agent_id}"
+        if full_path.exists():
+            # recursive delete
+            shutil.rmtree(full_path)

core/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`)
+# from llama_index.llms import Anthropic
+# # set Anthropic key
+# os.environ["ANTHROPIC_API_KEY"] = st.secrets.anthropic_key
+# BUILDER_LLM = Anthropic()

core/callback_manager.py

@@ -0,0 +1,70 @@
+"""Streaming callback manager."""
+from llama_index.callbacks.base_handler import BaseCallbackHandler
+from llama_index.callbacks.schema import CBEventType
+
+from typing import Optional, Dict, Any, List, Callable
+
+STORAGE_DIR = "./storage"  # directory to cache the generated index
+DATA_DIR = "./data"  # directory containing the documents to index
+
+
+class StreamlitFunctionsCallbackHandler(BaseCallbackHandler):
+    """Callback handler that outputs streamlit components given events."""
+
+    def __init__(self, msg_handler: Callable[[str], Any]) -> None:
+        """Initialize the base callback handler."""
+        self.msg_handler = msg_handler
+        super().__init__([], [])
+
+    def on_event_start(
+        self,
+        event_type: CBEventType,
+        payload: Optional[Dict[str, Any]] = None,
+        event_id: str = "",
+        parent_id: str = "",
+        **kwargs: Any,
+    ) -> str:
+        """Run when an event starts and return id of event."""
+        if event_type == CBEventType.FUNCTION_CALL:
+            if payload is None:
+                raise ValueError("Payload cannot be None")
+            arguments_str = payload["function_call"]
+            tool_str = payload["tool"].name
+            print_str = f"Calling function: {tool_str} with args: {arguments_str}\n\n"
+            self.msg_handler(print_str)
+        else:
+            pass
+        return event_id
+
+    def on_event_end(
+        self,
+        event_type: CBEventType,
+        payload: Optional[Dict[str, Any]] = None,
+        event_id: str = "",
+        **kwargs: Any,
+    ) -> None:
+        """Run when an event ends."""
+        pass
+        # TODO: currently we don't need to do anything here
+        # if event_type == CBEventType.FUNCTION_CALL:
+        #     response = payload["function_call_response"]
+        #     # Add this to queue
+        #     print_str = (
+        #         f"\n\nGot output: {response}\n"
+        #         "========================\n\n"
+        #     )
+        # elif event_type == CBEventType.AGENT_STEP:
+        #     # put response into queue
+        #     self._queue.put(payload["response"])
+
+    def start_trace(self, trace_id: Optional[str] = None) -> None:
+        """Run when an overall trace is launched."""
+        pass
+
+    def end_trace(
+        self,
+        trace_id: Optional[str] = None,
+        trace_map: Optional[Dict[str, List[str]]] = None,
+    ) -> None:
+        """Run when an overall trace is exited."""
+        pass

core/constants.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+AGENT_CACHE_DIR = Path(__file__).parent.parent / "cache" / "agents"
+MESSAGES_CACHE_DIR = Path(__file__).parent.parent / "cache" / "messages"

core/param_cache.py

@@ -0,0 +1,156 @@
+"""Param cache."""
+
+from pydantic import BaseModel, Field
+from llama_index import (
+    VectorStoreIndex,
+    StorageContext,
+    load_index_from_storage,
+)
+from typing import List, cast, Optional
+from llama_index.chat_engine.types import BaseChatEngine
+from pathlib import Path
+import json
+import uuid
+from core.utils import (
+    load_data,
+    get_tool_objects,
+    construct_agent,
+    RAGParams,
+    construct_mm_agent,
+)
+
+
+class ParamCache(BaseModel):
+    """Cache for RAG agent builder.
+
+    Created a wrapper class around a dict in case we wanted to more explicitly
+    type different items in the cache.
+
+    """
+
+    # arbitrary types
+    class Config:
+        arbitrary_types_allowed = True
+
+    # system prompt
+    system_prompt: Optional[str] = Field(
+        default=None, description="System prompt for RAG agent."
+    )
+    # data
+    file_names: List[str] = Field(
+        default_factory=list, description="File names as data source (if specified)"
+    )
+    urls: List[str] = Field(
+        default_factory=list, description="URLs as data source (if specified)"
+    )
+    directory: Optional[str] = Field(
+        default=None, description="Directory as data source (if specified)"
+    )
+
+    docs: List = Field(default_factory=list, description="Documents for RAG agent.")
+    # tools
+    tools: List = Field(
+        default_factory=list, description="Additional tools for RAG agent (e.g. web)"
+    )
+    # RAG params
+    rag_params: RAGParams = Field(
+        default_factory=RAGParams, description="RAG parameters for RAG agent."
+    )
+
+    # agent params
+    builder_type: str = Field(
+        default="default", description="Builder type (default, multimodal)."
+    )
+    vector_index: Optional[VectorStoreIndex] = Field(
+        default=None, description="Vector index for RAG agent."
+    )
+    agent_id: str = Field(
+        default_factory=lambda: f"Agent_{str(uuid.uuid4())}",
+        description="Agent ID for RAG agent.",
+    )
+    agent: Optional[BaseChatEngine] = Field(default=None, description="RAG agent.")
+
+    def save_to_disk(self, save_dir: str) -> None:
+        """Save cache to disk."""
+        # NOTE: more complex than just calling dict() because we want to
+        # only store serializable fields and be space-efficient
+
+        dict_to_serialize = {
+            "system_prompt": self.system_prompt,
+            "file_names": self.file_names,
+            "urls": self.urls,
+            "directory": self.directory,
+            # TODO: figure out tools
+            "tools": self.tools,
+            "rag_params": self.rag_params.dict(),
+            "builder_type": self.builder_type,
+            "agent_id": self.agent_id,
+        }
+        # store the vector store within the agent
+        if self.vector_index is None:
+            raise ValueError("Must specify vector index in order to save.")
+        self.vector_index.storage_context.persist(Path(save_dir) / "storage")
+
+        # if save_path directories don't exist, create it
+        if not Path(save_dir).exists():
+            Path(save_dir).mkdir(parents=True)
+        with open(Path(save_dir) / "cache.json", "w") as f:
+            json.dump(dict_to_serialize, f)
+
+    @classmethod
+    def load_from_disk(
+        cls,
+        save_dir: str,
+    ) -> "ParamCache":
+        """Load cache from disk."""
+        with open(Path(save_dir) / "cache.json", "r") as f:
+            cache_dict = json.load(f)
+
+        storage_context = StorageContext.from_defaults(
+            persist_dir=str(Path(save_dir) / "storage")
+        )
+        if cache_dict["builder_type"] == "multimodal":
+            from llama_index.indices.multi_modal.base import MultiModalVectorStoreIndex
+
+            vector_index: VectorStoreIndex = cast(
+                MultiModalVectorStoreIndex, load_index_from_storage(storage_context)
+            )
+        else:
+            vector_index = cast(
+                VectorStoreIndex, load_index_from_storage(storage_context)
+            )
+
+        # replace rag params with RAGParams object
+        cache_dict["rag_params"] = RAGParams(**cache_dict["rag_params"])
+
+        # add in the missing fields
+        # load docs
+        cache_dict["docs"] = load_data(
+            file_names=cache_dict["file_names"],
+            urls=cache_dict["urls"],
+            directory=cache_dict["directory"],
+        )
+        # load agent from index
+        additional_tools = get_tool_objects(cache_dict["tools"])
+
+        if cache_dict["builder_type"] == "multimodal":
+            vector_index = cast(MultiModalVectorStoreIndex, vector_index)
+            agent, _ = construct_mm_agent(
+                cache_dict["system_prompt"],
+                cache_dict["rag_params"],
+                cache_dict["docs"],
+                mm_vector_index=vector_index,
+            )
+        else:
+            agent, _ = construct_agent(
+                cache_dict["system_prompt"],
+                cache_dict["rag_params"],
+                cache_dict["docs"],
+                vector_index=vector_index,
+                additional_tools=additional_tools,
+                # TODO: figure out tools
+            )
+        cache_dict["vector_index"] = vector_index
+        cache_dict["agent"] = agent
+
+        return cls(**cache_dict)

core/utils.py

@@ -0,0 +1,480 @@
+"""Utils."""
+
+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
+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,
+    Document,
+)
+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
+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
+import streamlit as st
+
+from llama_index.callbacks import CallbackManager, trace_method
+from core.callback_manager import StreamlitFunctionsCallbackHandler
+from llama_index.schema import ImageNode, NodeWithScore
+
+### BETA: Multi-modal
+from llama_index.indices.multi_modal.base import MultiModalVectorStoreIndex
+from llama_index.multi_modal_llms.openai import OpenAIMultiModal
+from llama_index.indices.multi_modal.retriever import (
+    MultiModalVectorIndexRetriever,
+)
+from llama_index.llms import ChatMessage
+from llama_index.query_engine.multi_modal import SimpleMultiModalQueryEngine
+from llama_index.chat_engine.types import (
+    AGENT_CHAT_RESPONSE_TYPE,
+    StreamingAgentChatResponse,
+    AgentChatResponse,
+)
+from llama_index.llms.base import ChatResponse
+from typing import Generator
+
+
+class RAGParams(BaseModel):
+    """RAG parameters.
+
+    Parameters used to configure a RAG pipeline.
+
+    """
+
+    include_summarization: bool = Field(
+        default=False,
+        description=(
+            "Whether to include summarization in the RAG pipeline. (only for GPT-4)"
+        ),
+    )
+    top_k: int = Field(
+        default=2, description="Number of documents to retrieve from vector store."
+    )
+    chunk_size: int = Field(default=1024, description="Chunk size for vector store.")
+    embed_model: str = Field(
+        default="default", description="Embedding model to use (default is OpenAI)"
+    )
+    llm: str = Field(
+        default="gpt-4-1106-preview", description="LLM to use for summarization."
+    )
+
+
+def _resolve_llm(llm_str: str) -> LLM:
+    """Resolve LLM."""
+    # TODO: make this less hardcoded with if-else statements
+    # see if there's a prefix
+    # - if there isn't, assume it's an OpenAI model
+    # - if there is, resolve it
+    tokens = llm_str.split(":")
+    if len(tokens) == 1:
+        os.environ["OPENAI_API_KEY"] = st.secrets.openai_key
+        llm: LLM = OpenAI(model=llm_str)
+    elif tokens[0] == "local":
+        llm = resolve_llm(llm_str)
+    elif tokens[0] == "openai":
+        os.environ["OPENAI_API_KEY"] = st.secrets.openai_key
+        llm = OpenAI(model=tokens[1])
+    elif tokens[0] == "anthropic":
+        os.environ["ANTHROPIC_API_KEY"] = st.secrets.anthropic_key
+        llm = Anthropic(model=tokens[1])
+    elif tokens[0] == "replicate":
+        os.environ["REPLICATE_API_KEY"] = st.secrets.replicate_key
+        llm = Replicate(model=tokens[1])
+    else:
+        raise ValueError(f"LLM {llm_str} not recognized.")
+    return llm
+
+
+def load_data(
+    file_names: Optional[List[str]] = None,
+    directory: Optional[str] = None,
+    urls: Optional[List[str]] = None,
+) -> List[Document]:
+    """Load data."""
+    file_names = file_names or []
+    directory = directory or ""
+    urls = urls or []
+
+    # get number depending on whether specified
+    num_specified = sum(1 for v in [file_names, urls, directory] if v)
+
+    if num_specified == 0:
+        raise ValueError("Must specify either file_names or urls or directory.")
+    elif num_specified > 1:
+        raise ValueError("Must specify only one of file_names or urls or directory.")
+    elif file_names:
+        reader = SimpleDirectoryReader(input_files=file_names)
+        docs = reader.load_data()
+    elif directory:
+        reader = SimpleDirectoryReader(input_dir=directory)
+        docs = reader.load_data()
+    elif urls:
+        from llama_hub.web.simple_web.base import SimpleWebPageReader
+
+        # use simple web page reader from llamahub
+        loader = SimpleWebPageReader()
+        docs = loader.load_data(urls=urls)
+    else:
+        raise ValueError("Must specify either file_names or urls or directory.")
+
+    return docs
+
+
+def load_agent(
+    tools: List,
+    llm: LLM,
+    system_prompt: str,
+    extra_kwargs: Optional[Dict] = None,
+    **kwargs: Any,
+) -> BaseChatEngine:
+    """Load agent."""
+    extra_kwargs = extra_kwargs or {}
+    if isinstance(llm, OpenAI) and is_function_calling_model(llm.model):
+        # TODO: use default msg handler
+        # TODO: separate this from agent_utils.py...
+        def _msg_handler(msg: str) -> None:
+            """Message handler."""
+            st.info(msg)
+            st.session_state.agent_messages.append(
+                {"role": "assistant", "content": msg, "msg_type": "info"}
+            )
+
+        # add streamlit callbacks (to inject events)
+        handler = StreamlitFunctionsCallbackHandler(_msg_handler)
+        callback_manager = CallbackManager([handler])
+        # get OpenAI Agent
+        agent: BaseChatEngine = OpenAIAgent.from_tools(
+            tools=tools,
+            llm=llm,
+            system_prompt=system_prompt,
+            **kwargs,
+            callback_manager=callback_manager,
+        )
+    else:
+        if "vector_index" not in extra_kwargs:
+            raise ValueError(
+                "Must pass in vector index for CondensePlusContextChatEngine."
+            )
+        vector_index = cast(VectorStoreIndex, extra_kwargs["vector_index"])
+        rag_params = cast(RAGParams, extra_kwargs["rag_params"])
+        # use condense + context chat engine
+        agent = CondensePlusContextChatEngine.from_defaults(
+            vector_index.as_retriever(similarity_top_k=rag_params.top_k),
+        )
+
+    return agent
+
+
+def load_meta_agent(
+    tools: List,
+    llm: LLM,
+    system_prompt: str,
+    extra_kwargs: Optional[Dict] = None,
+    **kwargs: Any,
+) -> BaseAgent:
+    """Load meta agent.
+
+    TODO: consolidate with load_agent.
+
+    The meta-agent *has* to perform tool-use.
+
+    """
+    extra_kwargs = extra_kwargs or {}
+    if isinstance(llm, OpenAI) and is_function_calling_model(llm.model):
+        # get OpenAI Agent
+
+        agent: BaseAgent = 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 + "\n" + REACT_CHAT_SYSTEM_HEADER,
+            ),
+            **kwargs,
+        )
+
+    return agent
+
+
+def construct_agent(
+    system_prompt: str,
+    rag_params: RAGParams,
+    docs: List[Document],
+    vector_index: Optional[VectorStoreIndex] = None,
+    additional_tools: Optional[List] = None,
+) -> Tuple[BaseChatEngine, Dict]:
+    """Construct agent from docs / parameters / indices."""
+    extra_info = {}
+    additional_tools = additional_tools or []
+
+    # first resolve llm and embedding model
+    embed_model = resolve_embed_model(rag_params.embed_model)
+    # llm = resolve_llm(rag_params.llm)
+    # TODO: use OpenAI for now
+    # llm = OpenAI(model=rag_params.llm)
+    llm = _resolve_llm(rag_params.llm)
+
+    # first let's index the data with the right parameters
+    service_context = ServiceContext.from_defaults(
+        chunk_size=rag_params.chunk_size,
+        llm=llm,
+        embed_model=embed_model,
+    )
+
+    if vector_index is None:
+        vector_index = VectorStoreIndex.from_documents(
+            docs, service_context=service_context
+        )
+    else:
+        pass
+
+    extra_info["vector_index"] = vector_index
+
+    vector_query_engine = vector_index.as_query_engine(
+        similarity_top_k=rag_params.top_k
+    )
+    all_tools = []
+    vector_tool = QueryEngineTool(
+        query_engine=vector_query_engine,
+        metadata=ToolMetadata(
+            name="vector_tool",
+            description=("Use this tool to answer any user question over any data."),
+        ),
+    )
+    all_tools.append(vector_tool)
+    if rag_params.include_summarization:
+        summary_index = SummaryIndex.from_documents(
+            docs, service_context=service_context
+        )
+        summary_query_engine = summary_index.as_query_engine()
+        summary_tool = QueryEngineTool(
+            query_engine=summary_query_engine,
+            metadata=ToolMetadata(
+                name="summary_tool",
+                description=(
+                    "Use this tool for any user questions that ask "
+                    "for a summarization of content"
+                ),
+            ),
+        )
+        all_tools.append(summary_tool)
+
+    # then we add tools
+    all_tools.extend(additional_tools)
+
+    # build agent
+    if system_prompt is None:
+        return "System prompt not set yet. Please set system prompt first."
+
+    agent = load_agent(
+        all_tools,
+        llm=llm,
+        system_prompt=system_prompt,
+        verbose=True,
+        extra_kwargs={"vector_index": vector_index, "rag_params": rag_params},
+    )
+    return agent, extra_info
+
+
+def get_web_agent_tool() -> QueryEngineTool:
+    """Get web agent tool.
+
+    Wrap with our load and search tool spec.
+
+    """
+    from llama_hub.tools.metaphor.base import MetaphorToolSpec
+
+    # TODO: set metaphor API key
+    metaphor_tool = MetaphorToolSpec(
+        api_key=st.secrets.metaphor_key,
+    )
+    metaphor_tool_list = metaphor_tool.to_tool_list()
+
+    # TODO: LoadAndSearch doesn't work yet
+    # The search_and_retrieve_documents tool is the third in the tool list,
+    # as seen above
+    # wrapped_retrieve = LoadAndSearchToolSpec.from_defaults(
+    #     metaphor_tool_list[2],
+    # )
+
+    # NOTE: requires openai right now
+    # We don't give the Agent our unwrapped retrieve document tools
+    # instead passing the wrapped tools
+    web_agent = OpenAIAgent.from_tools(
+        # [*wrapped_retrieve.to_tool_list(), metaphor_tool_list[4]],
+        metaphor_tool_list,
+        llm=BUILDER_LLM,
+        verbose=True,
+    )
+
+    # return agent as a tool
+    # TODO: tune description
+    web_agent_tool = QueryEngineTool.from_defaults(
+        web_agent,
+        name="web_agent",
+        description="""
+            This agent can answer questions by searching the web. \
+Use this tool if the answer is ONLY likely to be found by searching \
+the internet, especially for queries about recent events.
+        """,
+    )
+
+    return web_agent_tool
+
+
+def get_tool_objects(tool_names: List[str]) -> List:
+    """Get tool objects from tool names."""
+    # construct additional tools
+    tool_objs = []
+    for tool_name in tool_names:
+        if tool_name == "web_search":
+            # build web agent
+            tool_objs.append(get_web_agent_tool())
+        else:
+            raise ValueError(f"Tool {tool_name} not recognized.")
+
+    return tool_objs
+
+
+class MultimodalChatEngine(BaseChatEngine):
+    """Multimodal chat engine.
+
+    This chat engine is a light wrapper around a query engine.
+    Offers no real 'chat' functionality, is a beta feature.
+
+    """
+
+    def __init__(self, mm_query_engine: SimpleMultiModalQueryEngine) -> None:
+        """Init params."""
+        self._mm_query_engine = mm_query_engine
+
+    def reset(self) -> None:
+        """Reset conversation state."""
+        pass
+
+    @property
+    def chat_history(self) -> List[ChatMessage]:
+        return []
+
+    @trace_method("chat")
+    def chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> AGENT_CHAT_RESPONSE_TYPE:
+        """Main chat interface."""
+        # just return the top-k results
+        response = self._mm_query_engine.query(message)
+        return AgentChatResponse(
+            response=str(response), source_nodes=response.source_nodes
+        )
+
+    @trace_method("chat")
+    def stream_chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> StreamingAgentChatResponse:
+        """Stream chat interface."""
+        response = self._mm_query_engine.query(message)
+
+        def _chat_stream(response: str) -> Generator[ChatResponse, None, None]:
+            yield ChatResponse(message=ChatMessage(role="assistant", content=response))
+
+        chat_stream = _chat_stream(str(response))
+        return StreamingAgentChatResponse(
+            chat_stream=chat_stream, source_nodes=response.source_nodes
+        )
+
+    @trace_method("chat")
+    async def achat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> AGENT_CHAT_RESPONSE_TYPE:
+        """Async version of main chat interface."""
+        response = await self._mm_query_engine.aquery(message)
+        return AgentChatResponse(
+            response=str(response), source_nodes=response.source_nodes
+        )
+
+    @trace_method("chat")
+    async def astream_chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> StreamingAgentChatResponse:
+        """Async version of main chat interface."""
+        return self.stream_chat(message, chat_history)
+
+
+def construct_mm_agent(
+    system_prompt: str,
+    rag_params: RAGParams,
+    docs: List[Document],
+    mm_vector_index: Optional[VectorStoreIndex] = None,
+    additional_tools: Optional[List] = None,
+) -> Tuple[BaseChatEngine, Dict]:
+    """Construct agent from docs / parameters / indices.
+
+    NOTE: system prompt isn't used right now
+
+    """
+    extra_info = {}
+    additional_tools = additional_tools or []
+
+    # first resolve llm and embedding model
+    embed_model = resolve_embed_model(rag_params.embed_model)
+    # TODO: use OpenAI for now
+    os.environ["OPENAI_API_KEY"] = st.secrets.openai_key
+    openai_mm_llm = OpenAIMultiModal(model="gpt-4-vision-preview", max_new_tokens=1500)
+
+    # first let's index the data with the right parameters
+    service_context = ServiceContext.from_defaults(
+        chunk_size=rag_params.chunk_size,
+        embed_model=embed_model,
+    )
+
+    if mm_vector_index is None:
+        mm_vector_index = MultiModalVectorStoreIndex.from_documents(
+            docs, service_context=service_context
+        )
+    else:
+        pass
+
+    mm_retriever = mm_vector_index.as_retriever(similarity_top_k=rag_params.top_k)
+    mm_query_engine = SimpleMultiModalQueryEngine(
+        cast(MultiModalVectorIndexRetriever, mm_retriever),
+        multi_modal_llm=openai_mm_llm,
+    )
+
+    extra_info["vector_index"] = mm_vector_index
+
+    # use condense + context chat engine
+    agent = MultimodalChatEngine(mm_query_engine)
+
+    return agent, extra_info
+
+
+def get_image_and_text_nodes(
+    nodes: List[NodeWithScore],
+) -> Tuple[List[NodeWithScore], List[NodeWithScore]]:
+    image_nodes = []
+    text_nodes = []
+    for res_node in nodes:
+        if isinstance(res_node.node, ImageNode):
+            image_nodes.append(res_node)
+        else:
+            text_nodes.append(res_node)
+    return image_nodes, text_nodes

pages/4_🤖_ChatDoctor.py

@@ -0,0 +1,126 @@
+"""Streamlit page showing builder config."""
+import streamlit as st
+from st_utils import add_sidebar, get_current_state
+from core.utils import get_image_and_text_nodes
+from llama_index.schema import MetadataMode
+from llama_index.chat_engine.types import AGENT_CHAT_RESPONSE_TYPE
+from typing import Dict, Optional
+import pandas as pd
+
+
+####################
+#### STREAMLIT #####
+####################
+
+
+st.set_page_config(
+    page_title="ChatDoctor: your virtual primary care physician assistant",
+    page_icon="🤖💬",
+    layout="centered",
+    #initial_sidebar_state="auto", #ggyimah set this to off
+    menu_items=None,
+)
+st.title("ChatDoctor: your virtual primary care physician assistant")
+#st.info(
+#    "Welcome!!! My name is ChatDoctor and I am trained to provide medical diagnoses and advice.",
+#    icon="ℹ️",
+#)
+
+
+
+current_state = get_current_state()
+add_sidebar()
+
+if (
+    "agent_messages" not in st.session_state.keys()
+):  # Initialize the chat messages history
+    st.session_state.agent_messages = [
+        {"role": "assistant", "content": "I am trained to provide medical diagnoses and advice. How may I help you, today?"}
+    ]
+
+
+def display_sources(response: AGENT_CHAT_RESPONSE_TYPE) -> None:
+    image_nodes, text_nodes = get_image_and_text_nodes(response.source_nodes)
+    if len(image_nodes) > 0 or len(text_nodes) > 0:
+        with st.expander("Sources"):
+            # get image nodes
+            if len(image_nodes) > 0:
+                st.subheader("Images")
+                for image_node in image_nodes:
+                    st.image(image_node.metadata["file_path"])
+
+            if len(text_nodes) > 0:
+                st.subheader("Text")
+                sources_df_list = []
+                for text_node in text_nodes:
+                    sources_df_list.append(
+                        {
+                            "ID": text_node.id_,
+                            "Text": text_node.node.get_content(
+                                metadata_mode=MetadataMode.ALL
+                            ),
+                        }
+                    )
+                sources_df = pd.DataFrame(sources_df_list)
+                st.dataframe(sources_df)
+
+
+def add_to_message_history(
+    role: str, content: str, extra: Optional[Dict] = None
+) -> None:
+    message = {"role": role, "content": str(content), "extra": extra}
+    st.session_state.agent_messages.append(message)  # Add response to message history
+
+
+def display_messages() -> None:
+    """Display messages."""
+    for message in st.session_state.agent_messages:  # Display the prior chat messages
+        with st.chat_message(message["role"]):
+            msg_type = message["msg_type"] if "msg_type" in message.keys() else "text"
+            if msg_type == "text":
+                st.write(message["content"])
+            elif msg_type == "info":
+                st.info(message["content"], icon="ℹ️")
+            else:
+                raise ValueError(f"Unknown message type: {msg_type}")
+
+            # display sources
+            if "extra" in message and isinstance(message["extra"], dict):
+                if "response" in message["extra"].keys():
+                    display_sources(message["extra"]["response"])
+
+
+# if agent is created, then we can chat with it
+if current_state.cache is not None and current_state.cache.agent is not None:
+    st.info(f"Viewing config for agent: {current_state.cache.agent_id}", icon="ℹ️")
+    agent = current_state.cache.agent
+
+    # display prior messages
+    display_messages()
+
+    # don't process selected for now
+    if prompt := st.chat_input(
+        "Your question"
+    ):  # Prompt for user input and save to chat history
+        add_to_message_history("user", prompt)
+        with st.chat_message("user"):
+            st.write(prompt)
+
+    # If last message is not from assistant, generate a new response
+    if st.session_state.agent_messages[-1]["role"] != "assistant":
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                response = agent.chat(str(prompt))
+                st.write(str(response))
+
+                # display sources
+                # Multi-modal: check if image nodes are present
+                display_sources(response)
+
+                add_to_message_history(
+                    "assistant", str(response), extra={"response": response}
+                )
+else:
+    st.info("In the side bar, select the ChatDoctor virtual agent (Agent_950acb55-056f-4324-957d-15e1c9b48695) to get started.\n")
+    st.info("Since this app is running on a free basic server, it could take from 2 to 10 minutes for the virtual agent to join you. \n Please be patient.")
+