From 12181ca1f68dcd5b37ff69cf45039adcbfc3dde0 Mon Sep 17 00:00:00 2001
From: Greg Hogue <gfjhogue@gmail.com>
Date: Wed, 26 Feb 2025 17:04:32 -0500
Subject: [PATCH 1/4] refactor `src/` while avoiding any real logic changes

---
 .github/actions/verify_imports.py             |   1 +
 bin/chat-chainlit.py                          |  17 +--
 bin/chat-repl                                 |  91 ------------
 bin/embeddings_manager                        |   4 +-
 src/__init__.py                               |   0
 src/{conversational_chain => agent}/graph.py  |  50 +++----
 src/agent/models.py                           |  68 +++++++++
 src/agent/tasks/rephrase.py                   |  31 +++++
 .../alliance/__init__.py}                     |   4 +-
 .../alliance/csv_generator.py}                |   0
 .../metadata_csv_loader.py                    |   0
 .../reactome/__init__.py}                     |   6 +-
 .../reactome/csv_generator.py}                |   6 +-
 .../reactome/neo4j_connector.py               |   0
 src/evaluation/evaluator.py                   |   6 +-
 src/evaluation/test_generator.py              |   2 +-
 src/retreival_chain.py                        | 130 ------------------
 .../chain.py => retrievers/rag_chain.py}      |  18 +--
 .../reactome/metadata_info.py                 |   0
 .../reactome/prompt.py}                       |  34 +----
 src/retrievers/reactome/rag.py                |  71 ++++++++++
 .../external_search/completeness_grader.py    |   2 +-
 src/{ => tools}/external_search/state.py      |   0
 .../external_search/tavily_wrapper.py         |   2 +-
 src/{ => tools}/external_search/workflow.py   |   6 +-
 25 files changed, 234 insertions(+), 315 deletions(-)
 delete mode 100644 bin/chat-repl
 delete mode 100644 src/__init__.py
 rename src/{conversational_chain => agent}/graph.py (80%)
 create mode 100644 src/agent/models.py
 create mode 100644 src/agent/tasks/rephrase.py
 rename src/{embeddings/alliance_generator.py => data_generation/alliance/__init__.py} (98%)
 mode change 100755 => 100644
 rename src/{csv_generator/alliance_generator.py => data_generation/alliance/csv_generator.py} (100%)
 rename src/{ => data_generation}/metadata_csv_loader.py (100%)
 rename src/{embeddings/reactome_generator.py => data_generation/reactome/__init__.py} (94%)
 mode change 100755 => 100644
 rename src/{csv_generator/reactome_generator.py => data_generation/reactome/csv_generator.py} (82%)
 rename src/{ => data_generation}/reactome/neo4j_connector.py (100%)
 delete mode 100644 src/retreival_chain.py
 rename src/{conversational_chain/chain.py => retrievers/rag_chain.py} (58%)
 rename src/{ => retrievers}/reactome/metadata_info.py (100%)
 rename src/{system_prompt/reactome_prompt.py => retrievers/reactome/prompt.py} (61%)
 create mode 100644 src/retrievers/reactome/rag.py
 rename src/{ => tools}/external_search/completeness_grader.py (97%)
 rename src/{ => tools}/external_search/state.py (100%)
 rename src/{ => tools}/external_search/tavily_wrapper.py (97%)
 rename src/{ => tools}/external_search/workflow.py (86%)

diff --git a/.github/actions/verify_imports.py b/.github/actions/verify_imports.py
index 6dea206..c1bf0a8 100644
--- a/.github/actions/verify_imports.py
+++ b/.github/actions/verify_imports.py
@@ -11,6 +11,7 @@
             "chat-chainlit.py",
             "chat-fastapi.py",
             "embeddings_manager",
+            "export_nologin_usage.py",
             "export_records.py",
         ],
     )
diff --git a/bin/chat-chainlit.py b/bin/chat-chainlit.py
index 6cebee2..f994c18 100644
--- a/bin/chat-chainlit.py
+++ b/bin/chat-chainlit.py
@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 import os
 from typing import Any
 
@@ -10,28 +8,17 @@
 from dotenv import load_dotenv
 from langchain_community.callbacks import OpenAICallbackHandler
 
-from conversational_chain.graph import RAGGraphWithMemory
-from retreival_chain import create_retrieval_chain
+from agent.graph import AgentGraph
 from util.chainlit_helpers import (is_feature_enabled, message_rate_limited,
                                    save_openai_metrics, static_messages,
                                    update_search_results)
 from util.config_yml import Config, TriggerEvent
-from util.embedding_environment import EmbeddingEnvironment
 from util.logging import logging
 
 load_dotenv()
 config: Config | None = Config.from_yaml()
 
-
-ENV = os.getenv("CHAT_ENV", "reactome")
-logging.info(f"Selected environment: {ENV}")
-
-llm_graph: RAGGraphWithMemory = create_retrieval_chain(
-    ENV,
-    EmbeddingEnvironment.get_dir(ENV),
-    hf_model=EmbeddingEnvironment.get_model(ENV),
-    oai_model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
-)
+llm_graph = AgentGraph()
 
 if os.getenv("POSTGRES_CHAINLIT_DB"):
     CHAINLIT_DB_URI = f"postgresql+psycopg://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_CHAINLIT_DB')}?sslmode=disable"
diff --git a/bin/chat-repl b/bin/chat-repl
deleted file mode 100644
index 3b72616..0000000
--- a/bin/chat-repl
+++ /dev/null
@@ -1,91 +0,0 @@
-#!/usr/bin/env python
-
-import argparse
-import os
-
-import pyfiglet
-from dotenv import load_dotenv
-
-from conversational_chain.chain import RAGChainWithMemory
-from retreival_chain import initialize_retrieval_chain
-from util.embedding_environment import EmbeddingEnvironment
-
-
-async def main() -> None:
-    load_dotenv()
-
-    parser = argparse.ArgumentParser(description="Reactome ChatBot")
-    parser.add_argument("--openai-key", help="API key for OpenAI")
-    parser.add_argument("--query", help="Query string to run")
-    parser.add_argument("--verbose", action="store_true", help="Enable verbose mode")
-    parser.add_argument(
-        "--ollama-model", help="Ollama language model to use (alternative to OpenAI)"
-    )
-    parser.add_argument(
-        "--ollama-url", default="http://localhost:11434", help="Ollama host url"
-    )
-    parser.add_argument(
-        "--hf-model",
-        help="HuggingFace sentence_transformers model (alternative to OpenAI)",
-    )
-    parser.add_argument(
-        "--hf-key",
-        help="API key for HuggingFaceHub",
-    )
-    parser.add_argument(
-        "--device",
-        default="cpu",
-        help="PyTorch device to use when running HuggingFace embeddings locally [cpu/cuda]",
-    )
-    args = parser.parse_args()
-
-    if args.openai_key:
-        os.environ["OPENAI_API_KEY"] = args.openai_key
-
-    if args.hf_key:
-        os.environ["HUGGINGFACEHUB_API_TOKEN"] = args.hf_key
-
-    env = os.getenv("CHAT_ENV", "reactome")
-    embeddings_directory = EmbeddingEnvironment.get_dir(env)
-
-    qa: RAGChainWithMemory = initialize_retrieval_chain(
-        env,
-        embeddings_directory,
-        True,
-        args.verbose,
-        args.ollama_model,
-        args.ollama_url,
-        args.hf_model,
-        args.device,
-    )
-    if args.query:
-        await print_results(qa, args.query, args.verbose)
-    else:
-        await interactive_mode(qa, args.verbose)
-
-
-async def interactive_mode(qa: RAGChainWithMemory, verbose: bool) -> None:
-    reactome_figlet = pyfiglet.figlet_format("React-to-me")
-    print(reactome_figlet)
-    print(
-        "Reactome Chatbot instructions: After each response you will have an opportunity to ask another questions. If you are done type enter instead of a question to exit."
-    )
-    while True:
-        query: str = input("\n\nUser Query:")
-        if not query:
-            break
-        print("\nResponse:")
-        await print_results(qa, query, verbose)
-
-
-async def print_results(
-    qa: RAGChainWithMemory, query: str, verbose: bool
-) -> None:
-    async for qa_result in qa.invoke(query):
-        pass
-
-
-if __name__ == "__main__":
-    import asyncio
-
-    asyncio.run(main())
diff --git a/bin/embeddings_manager b/bin/embeddings_manager
index d74e253..893cbb8 100755
--- a/bin/embeddings_manager
+++ b/bin/embeddings_manager
@@ -12,8 +12,8 @@ import boto3
 from botocore import UNSIGNED
 from botocore.client import Config
 
-from embeddings.alliance_generator import generate_alliance_embeddings
-from embeddings.reactome_generator import generate_reactome_embeddings
+from data_generation.alliance import generate_alliance_embeddings
+from data_generation.reactome import generate_reactome_embeddings
 from util.embedding_environment import EM_ARCHIVE, EmbeddingEnvironment
 
 S3_BUCKET = "download.reactome.org"
diff --git a/src/__init__.py b/src/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/conversational_chain/graph.py b/src/agent/graph.py
similarity index 80%
rename from src/conversational_chain/graph.py
rename to src/agent/graph.py
index e1e41de..b8f8887 100644
--- a/src/conversational_chain/graph.py
+++ b/src/agent/graph.py
@@ -4,9 +4,9 @@
 
 from langchain_core.callbacks.base import Callbacks
 from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
-from langchain_core.retrievers import BaseRetriever
 from langchain_core.runnables import Runnable, RunnableConfig
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.memory import MemorySaver
@@ -16,18 +16,15 @@
 from psycopg import AsyncConnection
 from psycopg_pool import AsyncConnectionPool
 
-from conversational_chain.chain import create_rag_chain, create_rephrase_chain
-from external_search.state import WebSearchResult
-from external_search.workflow import create_search_workflow
+from agent.models import get_embedding, get_llm
+from agent.tasks.rephrase import create_rephrase_chain
+from retrievers.reactome.rag import create_reactome_rag
+from tools.external_search.state import WebSearchResult
+from tools.external_search.workflow import create_search_workflow
 from util.logging import logging
 
 LANGGRAPH_DB_URI = f"postgresql://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_LANGGRAPH_DB')}?sslmode=disable"
 
-connection_kwargs = {
-    "autocommit": True,
-    "prepare_threshold": 0,
-}
-
 if not os.getenv("POSTGRES_LANGGRAPH_DB"):
     logging.warning("POSTGRES_LANGGRAPH_DB undefined; falling back to MemorySaver.")
 
@@ -36,26 +33,30 @@ class AdditionalContent(TypedDict):
     search_results: list[WebSearchResult]
 
 
-class ChatState(TypedDict):
+class AgentState(TypedDict):
     user_input: str  # User input text
     rephrased_input: str  # LLM-generated query from user input
     chat_history: Annotated[list[BaseMessage], add_messages]
     context: list[Document]
     answer: str  # primary LLM response that is streamed to the user
-    additional_content: (
-        AdditionalContent  # additional content to send after graph completes
-    )
+    additional_content: AdditionalContent  # sends on graph completion
+
 
+class AgentGraph:
+    def __init__(self) -> None:
+        # Get base models
+        llm: BaseChatModel = get_llm("openai", "gpt-4o-mini")
+        embedding: Embeddings = get_embedding("openai", "text-embedding-3-large")
 
-class RAGGraphWithMemory:
-    def __init__(self, retriever: BaseRetriever, llm: BaseChatModel) -> None:
-        # Set up runnables
-        self.rag_chain: Runnable = create_rag_chain(llm, retriever)
+        # Create runnables (tasks & tools)
+        self.reactome_rag: Runnable = create_reactome_rag(
+            llm, embedding, streaming=True
+        )
         self.rephrase_chain: Runnable = create_rephrase_chain(llm)
         self.search_workflow: CompiledStateGraph = create_search_workflow(llm)
 
         # Create graph
-        state_graph: StateGraph = StateGraph(ChatState)
+        state_graph = StateGraph(AgentState)
         # Set up nodes
         state_graph.add_node("preprocess", self.preprocess)
         state_graph.add_node("model", self.call_model)
@@ -88,7 +89,10 @@ async def create_checkpointer(self) -> BaseCheckpointSaver[str]:
             max_size=20,
             open=False,
             timeout=30,
-            kwargs=connection_kwargs,
+            kwargs={
+                "autocommit": True,
+                "prepare_threshold": 0,
+            },
         )
         await self.pool.open()
         checkpointer = AsyncPostgresSaver(self.pool)
@@ -100,15 +104,15 @@ async def close_pool(self) -> None:
             await self.pool.close()
 
     async def preprocess(
-        self, state: ChatState, config: RunnableConfig
+        self, state: AgentState, config: RunnableConfig
     ) -> dict[str, str]:
         query: str = await self.rephrase_chain.ainvoke(state, config)
         return {"rephrased_input": query}
 
     async def call_model(
-        self, state: ChatState, config: RunnableConfig
+        self, state: AgentState, config: RunnableConfig
     ) -> dict[str, Any]:
-        result: dict[str, Any] = await self.rag_chain.ainvoke(
+        result: dict[str, Any] = await self.reactome_rag.ainvoke(
             {
                 "input": state["rephrased_input"],
                 "user_input": state["user_input"],
@@ -126,7 +130,7 @@ async def call_model(
         }
 
     async def postprocess(
-        self, state: ChatState, config: RunnableConfig
+        self, state: AgentState, config: RunnableConfig
     ) -> dict[str, dict[str, list[WebSearchResult]]]:
         search_results: list[WebSearchResult] = []
         if config["configurable"]["enable_postprocess"]:
diff --git a/src/agent/models.py b/src/agent/models.py
new file mode 100644
index 0000000..01b324c
--- /dev/null
+++ b/src/agent/models.py
@@ -0,0 +1,68 @@
+from typing import Literal
+
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_huggingface import (HuggingFaceEmbeddings,
+                                   HuggingFaceEndpointEmbeddings)
+from langchain_ollama.chat_models import ChatOllama
+from langchain_openai.chat_models.base import ChatOpenAI
+from langchain_openai.embeddings import OpenAIEmbeddings
+
+
+def get_embedding(
+    provider: (
+        Literal[
+            "openai",
+            "huggingfacehub",
+            "huggingfacelocal",
+        ]
+        | str
+    ),
+    model: str | None = None,
+    *,
+    device: str | None = "cpu",
+) -> Embeddings:
+    if model is None:
+        provider, model = provider.split("/", 1)
+    if provider == "openai":
+        return OpenAIEmbeddings(model=model)
+    elif provider == "huggingfacehub":
+        return HuggingFaceEndpointEmbeddings(model=model)
+    elif provider == "huggingfacelocal":
+        return HuggingFaceEmbeddings(
+            model_name=model,
+            model_kwargs={"device": device, "trust_remote_code": True},
+            encode_kwargs={"batch_size": 12, "normalize_embeddings": False},
+        )
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+
+
+def get_llm(
+    provider: (
+        Literal[
+            "openai",
+            "ollama",
+        ]
+        | str
+    ),
+    model: str | None = None,
+    *,
+    base_url: str | None = None,
+) -> BaseChatModel:
+    if model is None:
+        provider, model = provider.split("/", 1)
+    if provider == "openai":
+        return ChatOpenAI(
+            model=model,
+            temperature=0.0,
+            base_url=base_url,
+        )
+    elif provider == "ollama":
+        return ChatOllama(
+            model=model,
+            temperature=0.0,
+            base_url=base_url,
+        )
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
diff --git a/src/agent/tasks/rephrase.py b/src/agent/tasks/rephrase.py
new file mode 100644
index 0000000..1851747
--- /dev/null
+++ b/src/agent/tasks/rephrase.py
@@ -0,0 +1,31 @@
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import Runnable
+
+contextualize_q_system_prompt = """
+You are an expert in question formulation with deep expertise in molecular biology and experience as a Reactome curator. Your task is to analyze the conversation history and the user’s latest query to fully understand their intent and what they seek to learn.
+If the user's question is not in English, reformulate the question and translate it to English, ensuring the meaning and intent are preserved.
+Reformulate the user’s question into a standalone version that retains its full meaning without requiring prior context. The reformulated question should be:
+    - Clear, concise, and precise
+    - Optimized for both vector search (semantic meaning) and case-sensitive keyword search
+    - Faithful to the user’s intent and scientific accuracy
+
+the returned question should always be in English.
+If the user’s question is already in English, self-contained and well-formed, return it as is.
+Do NOT answer the question or provide explanations.
+"""
+
+contextualize_q_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", contextualize_q_system_prompt),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{user_input}"),
+    ]
+)
+
+
+def create_rephrase_chain(llm: BaseChatModel) -> Runnable:
+    return (contextualize_q_prompt | llm | StrOutputParser()).with_config(
+        run_name="rephrase_question"
+    )
diff --git a/src/embeddings/alliance_generator.py b/src/data_generation/alliance/__init__.py
old mode 100755
new mode 100644
similarity index 98%
rename from src/embeddings/alliance_generator.py
rename to src/data_generation/alliance/__init__.py
index 18df90b..b0d2490
--- a/src/embeddings/alliance_generator.py
+++ b/src/data_generation/alliance/__init__.py
@@ -9,8 +9,8 @@
                                    HuggingFaceEndpointEmbeddings)
 from langchain_openai import OpenAIEmbeddings
 
-from csv_generator.alliance_generator import generate_all_csvs
-from metadata_csv_loader import MetaDataCSVLoader
+from data_generation.alliance.csv_generator import generate_all_csvs
+from data_generation.metadata_csv_loader import MetaDataCSVLoader
 
 
 def get_release_version() -> str:
diff --git a/src/csv_generator/alliance_generator.py b/src/data_generation/alliance/csv_generator.py
similarity index 100%
rename from src/csv_generator/alliance_generator.py
rename to src/data_generation/alliance/csv_generator.py
diff --git a/src/metadata_csv_loader.py b/src/data_generation/metadata_csv_loader.py
similarity index 100%
rename from src/metadata_csv_loader.py
rename to src/data_generation/metadata_csv_loader.py
diff --git a/src/embeddings/reactome_generator.py b/src/data_generation/reactome/__init__.py
old mode 100755
new mode 100644
similarity index 94%
rename from src/embeddings/reactome_generator.py
rename to src/data_generation/reactome/__init__.py
index 2cbe53f..fd73ec6
--- a/src/embeddings/reactome_generator.py
+++ b/src/data_generation/reactome/__init__.py
@@ -8,9 +8,9 @@
                                    HuggingFaceEndpointEmbeddings)
 from langchain_openai import OpenAIEmbeddings
 
-from csv_generator.reactome_generator import generate_all_csvs
-from metadata_csv_loader import MetaDataCSVLoader
-from reactome.neo4j_connector import Neo4jConnector
+from data_generation.metadata_csv_loader import MetaDataCSVLoader
+from data_generation.reactome.csv_generator import generate_all_csvs
+from data_generation.reactome.neo4j_connector import Neo4jConnector
 
 
 def upload_to_chromadb(
diff --git a/src/csv_generator/reactome_generator.py b/src/data_generation/reactome/csv_generator.py
similarity index 82%
rename from src/csv_generator/reactome_generator.py
rename to src/data_generation/reactome/csv_generator.py
index b283f61..15e3689 100644
--- a/src/csv_generator/reactome_generator.py
+++ b/src/data_generation/reactome/csv_generator.py
@@ -3,8 +3,10 @@
 
 import pandas as pd
 
-from reactome.neo4j_connector import (Neo4jConnector, Neo4jDict, get_complexes,
-                                      get_ewas, get_reactions, get_summations)
+from data_generation.reactome.neo4j_connector import (Neo4jConnector,
+                                                      Neo4jDict, get_complexes,
+                                                      get_ewas, get_reactions,
+                                                      get_summations)
 
 CSV_GENERATION_MAP: dict[str, Callable[[Neo4jConnector], Neo4jDict]] = {
     "reactions.csv": get_reactions,
diff --git a/src/reactome/neo4j_connector.py b/src/data_generation/reactome/neo4j_connector.py
similarity index 100%
rename from src/reactome/neo4j_connector.py
rename to src/data_generation/reactome/neo4j_connector.py
diff --git a/src/evaluation/evaluator.py b/src/evaluation/evaluator.py
index 45b5abe..c82c6e8 100644
--- a/src/evaluation/evaluator.py
+++ b/src/evaluation/evaluator.py
@@ -14,8 +14,9 @@
 from ragas.metrics import (answer_relevancy, context_recall,
                            context_utilization, faithfulness)
 
-from conversational_chain.chain import create_rag_chain
-from reactome.metadata_info import descriptions_info, field_info
+from retrievers.rag_chain import create_rag_chain
+from retrievers.reactome.metadata_info import descriptions_info, field_info
+from retrievers.reactome.prompt import qa_prompt
 
 
 def parse_arguments():
@@ -98,6 +99,7 @@ def initialize_rag_chain_with_memory(embeddings_directory, model_name, rag_type)
     qa = create_rag_chain(
         retriever=reactome_retriever,
         llm=llm,
+        qa_prompt=qa_prompt,
     )
     return qa
 
diff --git a/src/evaluation/test_generator.py b/src/evaluation/test_generator.py
index a1c5945..b7ea836 100644
--- a/src/evaluation/test_generator.py
+++ b/src/evaluation/test_generator.py
@@ -5,7 +5,7 @@
 from dotenv import load_dotenv
 from langchain_community.document_loaders import DirectoryLoader
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from ragas.testset.generator import TestsetGenerator
+from ragas.testset.synthesizers.generate import TestsetGenerator
 
 
 def parse_arguments():
diff --git a/src/retreival_chain.py b/src/retreival_chain.py
deleted file mode 100644
index c568c42..0000000
--- a/src/retreival_chain.py
+++ /dev/null
@@ -1,130 +0,0 @@
-import os
-from pathlib import Path
-from typing import Callable, Optional
-
-import chromadb.config
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-from langchain.retrievers import EnsembleRetriever
-from langchain.retrievers.merger_retriever import MergerRetriever
-from langchain.retrievers.self_query.base import SelfQueryRetriever
-from langchain_chroma.vectorstores import Chroma
-from langchain_community.document_loaders.csv_loader import CSVLoader
-from langchain_community.retrievers import BM25Retriever
-from langchain_core.embeddings import Embeddings
-from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.retrievers import BaseRetriever
-from langchain_huggingface import (HuggingFaceEmbeddings,
-                                   HuggingFaceEndpointEmbeddings)
-from langchain_ollama.chat_models import ChatOllama
-from langchain_openai.chat_models.base import ChatOpenAI
-from langchain_openai.embeddings import OpenAIEmbeddings
-
-from conversational_chain.graph import RAGGraphWithMemory
-from reactome.metadata_info import descriptions_info, field_info
-
-chroma_settings = chromadb.config.Settings(anonymized_telemetry=False)
-
-
-def list_chroma_subdirectories(directory: Path) -> list[str]:
-    subdirectories = list(
-        chroma_file.parent.name for chroma_file in directory.glob("*/chroma.sqlite3")
-    )
-    return subdirectories
-
-
-def get_embedding(
-    hf_model: Optional[str] = None, device: str = "cpu"
-) -> Callable[[], Embeddings]:
-    if hf_model is None:
-        return OpenAIEmbeddings
-    elif hf_model.startswith("openai/text-embedding-"):
-        return lambda: OpenAIEmbeddings(model=hf_model[len("openai/") :])
-    elif "HUGGINGFACEHUB_API_TOKEN" in os.environ:
-        return lambda: HuggingFaceEndpointEmbeddings(
-            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
-            model=hf_model,
-        )
-    else:
-        return lambda: HuggingFaceEmbeddings(
-            model_name=hf_model,
-            model_kwargs={"device": device, "trust_remote_code": True},
-            encode_kwargs={"batch_size": 12, "normalize_embeddings": False},
-        )
-
-
-def create_retrieval_chain(
-    env: str,
-    embeddings_directory: Path,
-    *,
-    commandline: bool = False,
-    verbose: bool = False,
-    ollama_model: Optional[str] = None,
-    ollama_url: str = "http://localhost:11434",
-    hf_model: Optional[str] = None,
-    oai_model: str = "gpt-4o-mini",
-    device: str = "cpu",
-) -> RAGGraphWithMemory:
-    callbacks: list[BaseCallbackHandler] = []
-    if commandline:
-        callbacks = [StreamingStdOutCallbackHandler()]
-
-    llm: BaseChatModel
-    if ollama_model is None:  # Use OpenAI when Ollama not specified
-        llm = ChatOpenAI(
-            temperature=0.0,
-            callbacks=callbacks,
-            verbose=verbose,
-            model=oai_model,
-        )
-    else:  # Otherwise use Ollama
-        llm = ChatOllama(
-            temperature=0.0,
-            callbacks=callbacks,
-            verbose=verbose,
-            model=ollama_model,
-            base_url=ollama_url,
-        )
-
-    # Get OpenAIEmbeddings (or HuggingFaceEmbeddings model if specified)
-    embedding_callable = get_embedding(hf_model, device)
-
-    # Adjusted type for retriever_list
-    retriever_list: list[BaseRetriever] = []
-    for subdirectory in list_chroma_subdirectories(embeddings_directory):
-        # set up BM25 retriever
-        csv_file_name = subdirectory + ".csv"
-        reactome_csvs_dir: Path = embeddings_directory / "csv_files"
-        loader = CSVLoader(file_path=reactome_csvs_dir / csv_file_name)
-        data = loader.load()
-        bm25_retriever = BM25Retriever.from_documents(data)
-        bm25_retriever.k = 10
-
-        # set up vectorstore SelfQuery retriever
-        embedding = embedding_callable()
-        vectordb = Chroma(
-            persist_directory=str(embeddings_directory / subdirectory),
-            embedding_function=embedding,
-            client_settings=chroma_settings,
-        )
-
-        selfq_retriever = SelfQueryRetriever.from_llm(
-            llm=llm,
-            vectorstore=vectordb,
-            document_contents=descriptions_info[subdirectory],
-            metadata_field_info=field_info[subdirectory],
-            search_kwargs={"k": 10},
-        )
-        rrf_retriever = EnsembleRetriever(
-            retrievers=[bm25_retriever, selfq_retriever], weights=[0.2, 0.8]
-        )
-        retriever_list.append(rrf_retriever)
-
-    reactome_retriever = MergerRetriever(retrievers=retriever_list)
-
-    qa = RAGGraphWithMemory(
-        retriever=reactome_retriever,
-        llm=llm,
-    )
-
-    return qa
diff --git a/src/conversational_chain/chain.py b/src/retrievers/rag_chain.py
similarity index 58%
rename from src/conversational_chain/chain.py
rename to src/retrievers/rag_chain.py
index 5b59510..3e5df8e 100644
--- a/src/conversational_chain/chain.py
+++ b/src/retrievers/rag_chain.py
@@ -1,23 +1,19 @@
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain.chains.retrieval import create_retrieval_chain
 from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.retrievers import BaseRetriever
 from langchain_core.runnables import Runnable
 
-from system_prompt.reactome_prompt import contextualize_q_prompt, qa_prompt
 
-
-def create_rephrase_chain(llm: BaseChatModel) -> Runnable:
-    return (contextualize_q_prompt | llm | StrOutputParser()).with_config(
-        run_name="rephrase_question"
-    )
-
-
-def create_rag_chain(llm: BaseChatModel, retriever: BaseRetriever) -> Runnable:
+def create_rag_chain(
+    llm: BaseChatModel,
+    retriever: BaseRetriever,
+    qa_prompt: ChatPromptTemplate,
+) -> Runnable:
     # Create the documents chain
     question_answer_chain: Runnable = create_stuff_documents_chain(
-        llm=llm.model_copy(update={"streaming": True}),
+        llm=llm,
         prompt=qa_prompt,
     )
 
diff --git a/src/reactome/metadata_info.py b/src/retrievers/reactome/metadata_info.py
similarity index 100%
rename from src/reactome/metadata_info.py
rename to src/retrievers/reactome/metadata_info.py
diff --git a/src/system_prompt/reactome_prompt.py b/src/retrievers/reactome/prompt.py
similarity index 61%
rename from src/system_prompt/reactome_prompt.py
rename to src/retrievers/reactome/prompt.py
index 0984708..1a004ed 100644
--- a/src/system_prompt/reactome_prompt.py
+++ b/src/retrievers/reactome/prompt.py
@@ -1,27 +1,5 @@
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 
-# Contextualize question prompt
-contextualize_q_system_prompt = """
-You are an expert in question formulation with deep expertise in molecular biology and experience as a Reactome curator. Your task is to analyze the conversation history and the user’s latest query to fully understand their intent and what they seek to learn.
-If the user's question is not in English, reformulate the question and translate it to English, ensuring the meaning and intent are preserved.
-Reformulate the user’s question into a standalone version that retains its full meaning without requiring prior context. The reformulated question should be:
-    - Clear, concise, and precise
-    - Optimized for both vector search (semantic meaning) and case-sensitive keyword search
-    - Faithful to the user’s intent and scientific accuracy
-
-the returned question should always be in English.
-If the user’s question is already in English, self-contained and well-formed, return it as is.
-Do NOT answer the question or provide explanations.
-"""
-
-contextualize_q_prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", contextualize_q_system_prompt),
-        MessagesPlaceholder(variable_name="chat_history"),
-        ("human", "{user_input}"),
-    ]
-)
-
 # Answer generation prompt
 qa_system_prompt = """
 You are an expert in molecular biology with access to the Reactome Knowledgebase.
@@ -30,17 +8,17 @@
 Always provide citations and links to the documents you obtained the information from.
 
 When providing answers, please adhere to the following guidelines:
-1. Provide answers **strictly based on the given context from the Reactome Knowledgebase**. Do **not** use or infer information from any external sources. 
+1. Provide answers **strictly based on the given context from the Reactome Knowledgebase**. Do **not** use or infer information from any external sources.
 2. If the answer cannot be derived from the context provided, do **not** answer the question; instead explain that the information is not currently available in Reactome.
 3. Answer the question comprehensively and accurately, providing useful background information based **only** on the context.
 4. keep track of **all** the sources that are directly used to derive the final answer, ensuring **every** piece of information in your response is **explicitly cited**.
-5. Create Citations for the sources used to generate the final asnwer according to the following: 
-     - For Reactome always format citations in the following format: <a href="url">*Source_Name*</a>, where *Source_Name* is the name of the retrieved document. 
-            Examples: 
+5. Create Citations for the sources used to generate the final asnwer according to the following:
+     - For Reactome always format citations in the following format: <a href="url">*Source_Name*</a>, where *Source_Name* is the name of the retrieved document.
+            Examples:
                 - <a href="https://reactome.org/content/detail/R-HSA-109581">Apoptosis</a>
                 - <a href="https://reactome.org/content/detail/R-HSA-1640170">Cell Cycle</a>
-                
-6. Always provide the citations you created in the format requested, in point-form at the end of the response paragraph, ensuring **every piece of information** provided in the final answer is cited. 
+
+6. Always provide the citations you created in the format requested, in point-form at the end of the response paragraph, ensuring **every piece of information** provided in the final answer is cited.
 7. Write in a conversational and engaging tone suitable for a chatbot.
 8. Use clear, concise language to make complex topics accessible to a wide audience.
 """
diff --git a/src/retrievers/reactome/rag.py b/src/retrievers/reactome/rag.py
new file mode 100644
index 0000000..701cbac
--- /dev/null
+++ b/src/retrievers/reactome/rag.py
@@ -0,0 +1,71 @@
+from pathlib import Path
+
+import chromadb.config
+from langchain.retrievers import EnsembleRetriever
+from langchain.retrievers.merger_retriever import MergerRetriever
+from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain_chroma.vectorstores import Chroma
+from langchain_community.document_loaders.csv_loader import CSVLoader
+from langchain_community.retrievers import BM25Retriever
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.runnables import Runnable
+
+from retrievers.rag_chain import create_rag_chain
+from retrievers.reactome.metadata_info import descriptions_info, field_info
+from retrievers.reactome.prompt import qa_prompt
+from util.embedding_environment import EmbeddingEnvironment
+
+chroma_settings = chromadb.config.Settings(anonymized_telemetry=False)
+
+
+def list_chroma_subdirectories(directory: Path) -> list[str]:
+    subdirectories = list(
+        chroma_file.parent.name for chroma_file in directory.glob("*/chroma.sqlite3")
+    )
+    return subdirectories
+
+
+def create_reactome_rag(
+    llm: BaseChatModel,
+    embedding: Embeddings,
+    embeddings_directory: Path = EmbeddingEnvironment.get_dir("reactome"),
+    *,
+    streaming: bool = False,
+) -> Runnable:
+    retriever_list: list[BaseRetriever] = []
+    for subdirectory in list_chroma_subdirectories(embeddings_directory):
+        # set up BM25 retriever
+        csv_file_name = subdirectory + ".csv"
+        reactome_csvs_dir: Path = embeddings_directory / "csv_files"
+        loader = CSVLoader(file_path=reactome_csvs_dir / csv_file_name)
+        data = loader.load()
+        bm25_retriever = BM25Retriever.from_documents(data)
+        bm25_retriever.k = 10
+
+        # set up vectorstore SelfQuery retriever
+        vectordb = Chroma(
+            persist_directory=str(embeddings_directory / subdirectory),
+            embedding_function=embedding,
+            client_settings=chroma_settings,
+        )
+
+        selfq_retriever = SelfQueryRetriever.from_llm(
+            llm=llm,
+            vectorstore=vectordb,
+            document_contents=descriptions_info[subdirectory],
+            metadata_field_info=field_info[subdirectory],
+            search_kwargs={"k": 10},
+        )
+        rrf_retriever = EnsembleRetriever(
+            retrievers=[bm25_retriever, selfq_retriever], weights=[0.2, 0.8]
+        )
+        retriever_list.append(rrf_retriever)
+
+    reactome_retriever = MergerRetriever(retrievers=retriever_list)
+
+    if streaming:
+        llm = llm.model_copy(update={"streaming": True})
+
+    return create_rag_chain(llm, reactome_retriever, qa_prompt)
diff --git a/src/external_search/completeness_grader.py b/src/tools/external_search/completeness_grader.py
similarity index 97%
rename from src/external_search/completeness_grader.py
rename to src/tools/external_search/completeness_grader.py
index 53baedc..9c33d2d 100644
--- a/src/external_search/completeness_grader.py
+++ b/src/tools/external_search/completeness_grader.py
@@ -3,7 +3,7 @@
 from langchain_core.runnables import Runnable, RunnableConfig
 from pydantic import BaseModel, Field
 
-from external_search.state import GraphState
+from tools.external_search.state import GraphState
 
 completeness_grader_message = """
 You are an expert grader with extensive knowledge in molecular biology and experience as a Reactome curator.
diff --git a/src/external_search/state.py b/src/tools/external_search/state.py
similarity index 100%
rename from src/external_search/state.py
rename to src/tools/external_search/state.py
diff --git a/src/external_search/tavily_wrapper.py b/src/tools/external_search/tavily_wrapper.py
similarity index 97%
rename from src/external_search/tavily_wrapper.py
rename to src/tools/external_search/tavily_wrapper.py
index 016d80f..d3a768a 100644
--- a/src/external_search/tavily_wrapper.py
+++ b/src/tools/external_search/tavily_wrapper.py
@@ -4,7 +4,7 @@
 
 from tavily import AsyncTavilyClient, MissingAPIKeyError
 
-from external_search.state import GraphState, WebSearchResult
+from tools.external_search.state import GraphState, WebSearchResult
 from util.logging import logging
 
 
diff --git a/src/external_search/workflow.py b/src/tools/external_search/workflow.py
similarity index 86%
rename from src/external_search/workflow.py
rename to src/tools/external_search/workflow.py
index 90c8f48..af746dd 100644
--- a/src/external_search/workflow.py
+++ b/src/tools/external_search/workflow.py
@@ -2,9 +2,9 @@
 from langgraph.graph import StateGraph
 from langgraph.graph.state import CompiledStateGraph
 
-from external_search.completeness_grader import CompletenessGrader
-from external_search.state import GraphState
-from external_search.tavily_wrapper import TavilyWrapper
+from tools.external_search.completeness_grader import CompletenessGrader
+from tools.external_search.state import GraphState
+from tools.external_search.tavily_wrapper import TavilyWrapper
 
 
 def decide_next_steps(state: GraphState) -> str:

From 7bd9d30953df80944d5f4a90719ff5764e3ae0f0 Mon Sep 17 00:00:00 2001
From: Greg Hogue <gfjhogue@gmail.com>
Date: Thu, 27 Feb 2025 18:11:27 -0500
Subject: [PATCH 2/4] add e2e chat profiles support [WIP]

---
 bin/chat-chainlit.py              |  14 ++--
 src/agent/graph.py                | 103 +++++-------------------------
 src/agent/profiles/__init__.py    |  33 ++++++++++
 src/agent/profiles/base.py        |  25 ++++++++
 src/agent/profiles/react_to_me.py |  92 ++++++++++++++++++++++++++
 5 files changed, 177 insertions(+), 90 deletions(-)
 create mode 100644 src/agent/profiles/__init__.py
 create mode 100644 src/agent/profiles/base.py
 create mode 100644 src/agent/profiles/react_to_me.py

diff --git a/bin/chat-chainlit.py b/bin/chat-chainlit.py
index f994c18..8b06e63 100644
--- a/bin/chat-chainlit.py
+++ b/bin/chat-chainlit.py
@@ -9,6 +9,7 @@
 from langchain_community.callbacks import OpenAICallbackHandler
 
 from agent.graph import AgentGraph
+from agent.profiles import CHAT_PROFILES
 from util.chainlit_helpers import (is_feature_enabled, message_rate_limited,
                                    save_openai_metrics, static_messages,
                                    update_search_results)
@@ -18,7 +19,9 @@
 load_dotenv()
 config: Config | None = Config.from_yaml()
 
-llm_graph = AgentGraph()
+profiles = ["React-to-Me"]
+
+llm_graph = AgentGraph(profiles)
 
 if os.getenv("POSTGRES_CHAINLIT_DB"):
     CHAINLIT_DB_URI = f"postgresql+psycopg://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_CHAINLIT_DB')}?sslmode=disable"
@@ -43,12 +46,13 @@ def oauth_callback(
 
 
 @cl.set_chat_profiles
-async def chat_profile() -> list[cl.ChatProfile]:
+async def chat_profiles() -> list[cl.ChatProfile]:
     return [
         cl.ChatProfile(
-            name="React-to-me",
-            markdown_description="An AI assistant specialized in exploring **Reactome** biological pathways and processes.",
+            name=CHAT_PROFILES[profile].name,
+            markdown_description=CHAT_PROFILES[profile].description,
         )
+        for profile in profiles
     ]
 
 
@@ -79,6 +83,7 @@ async def main(message: cl.Message) -> None:
     message_count: int = cl.user_session.get("message_count", 0) + 1
     cl.user_session.set("message_count", message_count)
 
+    chat_profile: str = cl.user_session.get("chat_profile")
     thread_id: str = cl.user_session.get("thread_id")
 
     chainlit_cb = cl.AsyncLangchainCallbackHandler(
@@ -90,6 +95,7 @@ async def main(message: cl.Message) -> None:
     enable_postprocess: bool = is_feature_enabled(config, "postprocessing")
     result: dict[str, Any] = await llm_graph.ainvoke(
         message.content,
+        chat_profile,
         callbacks=[chainlit_cb, openai_cb],
         thread_id=thread_id,
         enable_postprocess=enable_postprocess,
diff --git a/src/agent/graph.py b/src/agent/graph.py
index b8f8887..6bba134 100644
--- a/src/agent/graph.py
+++ b/src/agent/graph.py
@@ -1,26 +1,20 @@
 import asyncio
 import os
-from typing import Annotated, Any, TypedDict
+from typing import Any
 
 from langchain_core.callbacks.base import Callbacks
-from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
-from langchain_core.runnables import Runnable, RunnableConfig
+from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
-from langgraph.graph.message import add_messages
 from langgraph.graph.state import CompiledStateGraph, StateGraph
 from psycopg import AsyncConnection
 from psycopg_pool import AsyncConnectionPool
 
 from agent.models import get_embedding, get_llm
-from agent.tasks.rephrase import create_rephrase_chain
-from retrievers.reactome.rag import create_reactome_rag
-from tools.external_search.state import WebSearchResult
-from tools.external_search.workflow import create_search_workflow
+from agent.profiles import create_profile_graphs
 from util.logging import logging
 
 LANGGRAPH_DB_URI = f"postgresql://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_LANGGRAPH_DB')}?sslmode=disable"
@@ -29,57 +23,33 @@
     logging.warning("POSTGRES_LANGGRAPH_DB undefined; falling back to MemorySaver.")
 
 
-class AdditionalContent(TypedDict):
-    search_results: list[WebSearchResult]
-
-
-class AgentState(TypedDict):
-    user_input: str  # User input text
-    rephrased_input: str  # LLM-generated query from user input
-    chat_history: Annotated[list[BaseMessage], add_messages]
-    context: list[Document]
-    answer: str  # primary LLM response that is streamed to the user
-    additional_content: AdditionalContent  # sends on graph completion
-
-
 class AgentGraph:
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        profiles: list[str],
+    ) -> None:
         # Get base models
         llm: BaseChatModel = get_llm("openai", "gpt-4o-mini")
         embedding: Embeddings = get_embedding("openai", "text-embedding-3-large")
 
-        # Create runnables (tasks & tools)
-        self.reactome_rag: Runnable = create_reactome_rag(
-            llm, embedding, streaming=True
+        self.uncompiled_graph: dict[str, StateGraph] = create_profile_graphs(
+            profiles, llm, embedding
         )
-        self.rephrase_chain: Runnable = create_rephrase_chain(llm)
-        self.search_workflow: CompiledStateGraph = create_search_workflow(llm)
-
-        # Create graph
-        state_graph = StateGraph(AgentState)
-        # Set up nodes
-        state_graph.add_node("preprocess", self.preprocess)
-        state_graph.add_node("model", self.call_model)
-        state_graph.add_node("postprocess", self.postprocess)
-        # Set up edges
-        state_graph.set_entry_point("preprocess")
-        state_graph.add_edge("preprocess", "model")
-        state_graph.add_edge("model", "postprocess")
-        state_graph.set_finish_point("postprocess")
-
-        self.uncompiled_graph: StateGraph = state_graph
 
         # The following are set asynchronously by calling initialize()
-        self.graph: CompiledStateGraph | None = None
+        self.graph: dict[str, CompiledStateGraph] | None = None
         self.pool: AsyncConnectionPool[AsyncConnection[dict[str, Any]]] | None = None
 
     def __del__(self) -> None:
         if self.pool:
             asyncio.run(self.close_pool())
 
-    async def initialize(self) -> CompiledStateGraph:
+    async def initialize(self) -> dict[str, CompiledStateGraph]:
         checkpointer: BaseCheckpointSaver[str] = await self.create_checkpointer()
-        return self.uncompiled_graph.compile(checkpointer=checkpointer)
+        return {
+            profile: graph.compile(checkpointer=checkpointer)
+            for profile, graph in self.uncompiled_graph.items()
+        }
 
     async def create_checkpointer(self) -> BaseCheckpointSaver[str]:
         if not os.getenv("POSTGRES_LANGGRAPH_DB"):
@@ -103,49 +73,10 @@ async def close_pool(self) -> None:
         if self.pool:
             await self.pool.close()
 
-    async def preprocess(
-        self, state: AgentState, config: RunnableConfig
-    ) -> dict[str, str]:
-        query: str = await self.rephrase_chain.ainvoke(state, config)
-        return {"rephrased_input": query}
-
-    async def call_model(
-        self, state: AgentState, config: RunnableConfig
-    ) -> dict[str, Any]:
-        result: dict[str, Any] = await self.reactome_rag.ainvoke(
-            {
-                "input": state["rephrased_input"],
-                "user_input": state["user_input"],
-                "chat_history": state["chat_history"],
-            },
-            config,
-        )
-        return {
-            "chat_history": [
-                HumanMessage(state["user_input"]),
-                AIMessage(result["answer"]),
-            ],
-            "context": result["context"],
-            "answer": result["answer"],
-        }
-
-    async def postprocess(
-        self, state: AgentState, config: RunnableConfig
-    ) -> dict[str, dict[str, list[WebSearchResult]]]:
-        search_results: list[WebSearchResult] = []
-        if config["configurable"]["enable_postprocess"]:
-            result: dict[str, Any] = await self.search_workflow.ainvoke(
-                {"question": state["rephrased_input"], "generation": state["answer"]},
-                config=RunnableConfig(callbacks=config["callbacks"]),
-            )
-            search_results = result["search_results"]
-        return {
-            "additional_content": {"search_results": search_results},
-        }
-
     async def ainvoke(
         self,
         user_input: str,
+        profile: str,
         *,
         callbacks: Callbacks,
         thread_id: str,
@@ -153,7 +84,7 @@ async def ainvoke(
     ) -> dict[str, Any]:
         if self.graph is None:
             self.graph = await self.initialize()
-        result: dict[str, Any] = await self.graph.ainvoke(
+        result: dict[str, Any] = await self.graph[profile].ainvoke(
             {"user_input": user_input},
             config=RunnableConfig(
                 callbacks=callbacks,
diff --git a/src/agent/profiles/__init__.py b/src/agent/profiles/__init__.py
new file mode 100644
index 0000000..db7d428
--- /dev/null
+++ b/src/agent/profiles/__init__.py
@@ -0,0 +1,33 @@
+from typing import Callable, NamedTuple
+
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models.chat_models import BaseChatModel
+from langgraph.graph.state import StateGraph
+
+from agent.profiles.react_to_me import create_reacttome_graph
+
+
+class Profile(NamedTuple):
+    name: str
+    description: str
+    graph_builder: Callable[[BaseChatModel, Embeddings], StateGraph]
+
+
+CHAT_PROFILES: dict[str, Profile] = {
+    "React-to-Me": Profile(
+        name="React-to-Me",
+        description="An AI assistant specialized in exploring **Reactome** biological pathways and processes.",
+        graph_builder=create_reacttome_graph,
+    ),
+}
+
+
+def create_profile_graphs(
+    profiles: list[str],
+    llm: BaseChatModel,
+    embedding: Embeddings,
+) -> dict[str, StateGraph]:
+    return {
+        profile: CHAT_PROFILES[profile].graph_builder(llm, embedding)
+        for profile in profiles
+    }
diff --git a/src/agent/profiles/base.py b/src/agent/profiles/base.py
new file mode 100644
index 0000000..32d1a3f
--- /dev/null
+++ b/src/agent/profiles/base.py
@@ -0,0 +1,25 @@
+from typing import Annotated, TypedDict
+
+from langchain_core.documents import Document
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
+
+from tools.external_search.state import WebSearchResult
+
+
+class AdditionalContent(TypedDict):
+    search_results: list[WebSearchResult]
+
+
+class BaseState(TypedDict):
+    # (Everything the Chainlit layer uses should be included here)
+
+    user_input: str  # User input text
+    chat_history: Annotated[list[BaseMessage], add_messages]
+    context: list[Document]
+    answer: str  # primary LLM response that is streamed to the user
+    additional_content: AdditionalContent  # sends on graph completion
+
+
+class BaseGraphBuilder:
+    pass  # NOTE: Anything that is common to all graph builders goes here
diff --git a/src/agent/profiles/react_to_me.py b/src/agent/profiles/react_to_me.py
new file mode 100644
index 0000000..7583f2f
--- /dev/null
+++ b/src/agent/profiles/react_to_me.py
@@ -0,0 +1,92 @@
+from typing import Any
+
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.runnables import Runnable, RunnableConfig
+from langgraph.graph.state import CompiledStateGraph, StateGraph
+
+from agent.profiles.base import BaseGraphBuilder, BaseState
+from agent.tasks.rephrase import create_rephrase_chain
+from retrievers.reactome.rag import create_reactome_rag
+from tools.external_search.state import WebSearchResult
+from tools.external_search.workflow import create_search_workflow
+
+
+class ReactToMeState(BaseState):
+    rephrased_input: str  # LLM-generated query from user input
+
+
+class ReactToMeGraphBuilder(BaseGraphBuilder):
+    def __init__(
+        self,
+        llm: BaseChatModel,
+        embedding: Embeddings,
+    ) -> None:
+        # Create runnables (tasks & tools)
+        self.reactome_rag: Runnable = create_reactome_rag(
+            llm, embedding, streaming=True
+        )
+        self.rephrase_chain: Runnable = create_rephrase_chain(llm)
+        self.search_workflow: CompiledStateGraph = create_search_workflow(llm)
+
+        # Create graph
+        state_graph = StateGraph(ReactToMeState)
+        # Set up nodes
+        state_graph.add_node("preprocess", self.preprocess)
+        state_graph.add_node("model", self.call_model)
+        state_graph.add_node("postprocess", self.postprocess)
+        # Set up edges
+        state_graph.set_entry_point("preprocess")
+        state_graph.add_edge("preprocess", "model")
+        state_graph.add_edge("model", "postprocess")
+        state_graph.set_finish_point("postprocess")
+
+        self.uncompiled_graph: StateGraph = state_graph
+
+    async def preprocess(
+        self, state: ReactToMeState, config: RunnableConfig
+    ) -> dict[str, str]:
+        query: str = await self.rephrase_chain.ainvoke(state, config)
+        return {"rephrased_input": query}
+
+    async def call_model(
+        self, state: ReactToMeState, config: RunnableConfig
+    ) -> dict[str, Any]:
+        result: dict[str, Any] = await self.reactome_rag.ainvoke(
+            {
+                "input": state["rephrased_input"],
+                "user_input": state["user_input"],
+                "chat_history": state["chat_history"],
+            },
+            config,
+        )
+        return {
+            "chat_history": [
+                HumanMessage(state["user_input"]),
+                AIMessage(result["answer"]),
+            ],
+            "context": result["context"],
+            "answer": result["answer"],
+        }
+
+    async def postprocess(
+        self, state: ReactToMeState, config: RunnableConfig
+    ) -> dict[str, dict[str, list[WebSearchResult]]]:
+        search_results: list[WebSearchResult] = []
+        if config["configurable"]["enable_postprocess"]:
+            result: dict[str, Any] = await self.search_workflow.ainvoke(
+                {"question": state["rephrased_input"], "generation": state["answer"]},
+                config=RunnableConfig(callbacks=config["callbacks"]),
+            )
+            search_results = result["search_results"]
+        return {
+            "additional_content": {"search_results": search_results},
+        }
+
+
+def create_reacttome_graph(
+    llm: BaseChatModel,
+    embedding: Embeddings,
+) -> StateGraph:
+    return ReactToMeGraphBuilder(llm, embedding).uncompiled_graph

From 2943e319123fe569c7f30891ea65264d5d27be66 Mon Sep 17 00:00:00 2001
From: Greg Hogue <gfjhogue@gmail.com>
Date: Fri, 28 Feb 2025 11:13:37 -0500
Subject: [PATCH 3/4] add chat profiles to config

---
 .config.schema.yaml             |  7 ++++++-
 bin/chat-chainlit.py            | 14 +++++++-------
 config_default.yml              |  3 +++
 src/agent/graph.py              |  6 +++---
 src/agent/profiles/__init__.py  | 20 +++++++++++++++-----
 src/util/config_yml/__init__.py |  2 ++
 6 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/.config.schema.yaml b/.config.schema.yaml
index dd5acd8..10f67f6 100644
--- a/.config.schema.yaml
+++ b/.config.schema.yaml
@@ -51,6 +51,11 @@ properties:
             - required: ["event"]
             - required: ["after_messages"]
       required: ["message", "trigger"]
+  profiles:
+    type: array
+    items:
+      type: string
+      enum: ["React-to-Me"]
   usage_limits:
     type: object
     properties:
@@ -73,4 +78,4 @@ properties:
               pattern: "^[0-9]+[smhdw]$"
           required: ["users", "max_messages", "interval"]
     required: ["message_rates"]
-required: ["features", "messages", "usage_limits"]
+required: ["features", "messages", "profiles", "usage_limits"]
diff --git a/bin/chat-chainlit.py b/bin/chat-chainlit.py
index 8b06e63..5519e8f 100644
--- a/bin/chat-chainlit.py
+++ b/bin/chat-chainlit.py
@@ -9,7 +9,7 @@
 from langchain_community.callbacks import OpenAICallbackHandler
 
 from agent.graph import AgentGraph
-from agent.profiles import CHAT_PROFILES
+from agent.profiles import ProfileName, get_chat_profiles
 from util.chainlit_helpers import (is_feature_enabled, message_rate_limited,
                                    save_openai_metrics, static_messages,
                                    update_search_results)
@@ -19,8 +19,7 @@
 load_dotenv()
 config: Config | None = Config.from_yaml()
 
-profiles = ["React-to-Me"]
-
+profiles: list[ProfileName] = config.profiles if config else [ProfileName.React_to_Me]
 llm_graph = AgentGraph(profiles)
 
 if os.getenv("POSTGRES_CHAINLIT_DB"):
@@ -49,10 +48,10 @@ def oauth_callback(
 async def chat_profiles() -> list[cl.ChatProfile]:
     return [
         cl.ChatProfile(
-            name=CHAT_PROFILES[profile].name,
-            markdown_description=CHAT_PROFILES[profile].description,
+            name=profile.name,
+            markdown_description=profile.description,
         )
-        for profile in profiles
+        for profile in get_chat_profiles(profiles)
     ]
 
 
@@ -83,7 +82,8 @@ async def main(message: cl.Message) -> None:
     message_count: int = cl.user_session.get("message_count", 0) + 1
     cl.user_session.set("message_count", message_count)
 
-    chat_profile: str = cl.user_session.get("chat_profile")
+    chat_profile = ProfileName(cl.user_session.get("chat_profile"))
+
     thread_id: str = cl.user_session.get("thread_id")
 
     chainlit_cb = cl.AsyncLangchainCallbackHandler(
diff --git a/config_default.yml b/config_default.yml
index 88a9c41..6831216 100644
--- a/config_default.yml
+++ b/config_default.yml
@@ -1,5 +1,8 @@
 # yaml-language-server: $schema=./.config.schema.yaml
 
+profiles:
+  - React-to-Me
+
 features:
   postprocessing:  # external web search feature
     enabled: true
diff --git a/src/agent/graph.py b/src/agent/graph.py
index 6bba134..6727b7f 100644
--- a/src/agent/graph.py
+++ b/src/agent/graph.py
@@ -14,7 +14,7 @@
 from psycopg_pool import AsyncConnectionPool
 
 from agent.models import get_embedding, get_llm
-from agent.profiles import create_profile_graphs
+from agent.profiles import ProfileName, create_profile_graphs
 from util.logging import logging
 
 LANGGRAPH_DB_URI = f"postgresql://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_LANGGRAPH_DB')}?sslmode=disable"
@@ -26,7 +26,7 @@
 class AgentGraph:
     def __init__(
         self,
-        profiles: list[str],
+        profiles: list[ProfileName],
     ) -> None:
         # Get base models
         llm: BaseChatModel = get_llm("openai", "gpt-4o-mini")
@@ -76,7 +76,7 @@ async def close_pool(self) -> None:
     async def ainvoke(
         self,
         user_input: str,
-        profile: str,
+        profile: ProfileName,
         *,
         callbacks: Callbacks,
         thread_id: str,
diff --git a/src/agent/profiles/__init__.py b/src/agent/profiles/__init__.py
index db7d428..eafabf6 100644
--- a/src/agent/profiles/__init__.py
+++ b/src/agent/profiles/__init__.py
@@ -1,3 +1,4 @@
+from enum import StrEnum
 from typing import Callable, NamedTuple
 
 from langchain_core.embeddings import Embeddings
@@ -7,15 +8,20 @@
 from agent.profiles.react_to_me import create_reacttome_graph
 
 
+class ProfileName(StrEnum):
+    # These should exactly match names in .config.schema.yaml
+    React_to_Me = "React-to-Me"
+
+
 class Profile(NamedTuple):
-    name: str
+    name: ProfileName
     description: str
     graph_builder: Callable[[BaseChatModel, Embeddings], StateGraph]
 
 
-CHAT_PROFILES: dict[str, Profile] = {
-    "React-to-Me": Profile(
-        name="React-to-Me",
+CHAT_PROFILES: dict[ProfileName, Profile] = {
+    ProfileName.React_to_Me: Profile(
+        name=ProfileName.React_to_Me,
         description="An AI assistant specialized in exploring **Reactome** biological pathways and processes.",
         graph_builder=create_reacttome_graph,
     ),
@@ -23,7 +29,7 @@ class Profile(NamedTuple):
 
 
 def create_profile_graphs(
-    profiles: list[str],
+    profiles: list[ProfileName],
     llm: BaseChatModel,
     embedding: Embeddings,
 ) -> dict[str, StateGraph]:
@@ -31,3 +37,7 @@ def create_profile_graphs(
         profile: CHAT_PROFILES[profile].graph_builder(llm, embedding)
         for profile in profiles
     }
+
+
+def get_chat_profiles(profiles: list[ProfileName]) -> list[Profile]:
+    return [CHAT_PROFILES[profile] for profile in profiles]
diff --git a/src/util/config_yml/__init__.py b/src/util/config_yml/__init__.py
index 623588d..e6d57e9 100644
--- a/src/util/config_yml/__init__.py
+++ b/src/util/config_yml/__init__.py
@@ -4,6 +4,7 @@
 import yaml
 from pydantic import BaseModel, ValidationError
 
+from agent.profiles import ProfileName
 from util.config_yml.features import Feature, Features
 from util.config_yml.messages import Message, TriggerEvent
 from util.config_yml.usage_limits import MessageRate, UsageLimits
@@ -17,6 +18,7 @@
 class Config(BaseModel):
     features: Features
     messages: dict[str, Message]
+    profiles: list[ProfileName]
     usage_limits: UsageLimits
 
     def get_feature(

From 3c2d6edd6de944fdc2daf6f945d50f4c3ca5f3f6 Mon Sep 17 00:00:00 2001
From: Greg Hogue <gfjhogue@gmail.com>
Date: Fri, 28 Feb 2025 12:07:46 -0500
Subject: [PATCH 4/4] use lowercase to avoid compatibility issues with old
 profile naming

---
 bin/chat-chainlit.py           | 4 ++--
 src/agent/graph.py             | 4 +++-
 src/agent/profiles/__init__.py | 8 ++++----
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/bin/chat-chainlit.py b/bin/chat-chainlit.py
index 5519e8f..9de6e9f 100644
--- a/bin/chat-chainlit.py
+++ b/bin/chat-chainlit.py
@@ -82,7 +82,7 @@ async def main(message: cl.Message) -> None:
     message_count: int = cl.user_session.get("message_count", 0) + 1
     cl.user_session.set("message_count", message_count)
 
-    chat_profile = ProfileName(cl.user_session.get("chat_profile"))
+    chat_profile: str = cl.user_session.get("chat_profile")
 
     thread_id: str = cl.user_session.get("thread_id")
 
@@ -95,7 +95,7 @@ async def main(message: cl.Message) -> None:
     enable_postprocess: bool = is_feature_enabled(config, "postprocessing")
     result: dict[str, Any] = await llm_graph.ainvoke(
         message.content,
-        chat_profile,
+        chat_profile.lower(),
         callbacks=[chainlit_cb, openai_cb],
         thread_id=thread_id,
         enable_postprocess=enable_postprocess,
diff --git a/src/agent/graph.py b/src/agent/graph.py
index 6727b7f..de311d0 100644
--- a/src/agent/graph.py
+++ b/src/agent/graph.py
@@ -76,7 +76,7 @@ async def close_pool(self) -> None:
     async def ainvoke(
         self,
         user_input: str,
-        profile: ProfileName,
+        profile: str,
         *,
         callbacks: Callbacks,
         thread_id: str,
@@ -84,6 +84,8 @@ async def ainvoke(
     ) -> dict[str, Any]:
         if self.graph is None:
             self.graph = await self.initialize()
+        if profile not in self.graph:
+            return {}
         result: dict[str, Any] = await self.graph[profile].ainvoke(
             {"user_input": user_input},
             config=RunnableConfig(
diff --git a/src/agent/profiles/__init__.py b/src/agent/profiles/__init__.py
index eafabf6..92e3787 100644
--- a/src/agent/profiles/__init__.py
+++ b/src/agent/profiles/__init__.py
@@ -19,8 +19,8 @@ class Profile(NamedTuple):
     graph_builder: Callable[[BaseChatModel, Embeddings], StateGraph]
 
 
-CHAT_PROFILES: dict[ProfileName, Profile] = {
-    ProfileName.React_to_Me: Profile(
+CHAT_PROFILES: dict[str, Profile] = {
+    ProfileName.React_to_Me.lower(): Profile(
         name=ProfileName.React_to_Me,
         description="An AI assistant specialized in exploring **Reactome** biological pathways and processes.",
         graph_builder=create_reacttome_graph,
@@ -35,9 +35,9 @@ def create_profile_graphs(
 ) -> dict[str, StateGraph]:
     return {
         profile: CHAT_PROFILES[profile].graph_builder(llm, embedding)
-        for profile in profiles
+        for profile in map(str.lower, profiles)
     }
 
 
 def get_chat_profiles(profiles: list[ProfileName]) -> list[Profile]:
-    return [CHAT_PROFILES[profile] for profile in profiles]
+    return [CHAT_PROFILES[profile.lower()] for profile in profiles]
